102/102 encoder params load with 0 missing and 0 unmapped. A synthetic voxelised
sphere shell (16,934 voxels) encodes to 56 latent voxels in 135ms on m3ultra, and the
latent comes out mean +0.05 / std 0.92 - the approximately unit-normal distribution a
KL-trained VAE should produce, which is decent evidence the graph and the sparse conv
path are right.
Architecture is inferred from tensor shapes, not constructor defaults: upstream
defaults latent_dim to 8 but the released weights say 32, and attn_mode/pe_mode
defaults are likewise overridden by the trained config. infer_config() reads it off
the checkpoint.
Also added SparseDownsample. Upstream's docstring says average pooling but the code
passes reduce='amax' - following the code.
Kernel orientation: tried latent statistics as a cheap discriminator and it does NOT
work. The flip is not a no-op (max delta 3.53) but both orientations give a plausible
near-unit-normal latent (std 0.919 vs 0.945). Recorded as a negative result; it needs
the decoder and reconstruction quality to settle.
All verified against torch, which IS available here - only spconv was not. So unlike
the submanifold conv these compare against upstream's real semantics rather than a
hand-written paraphrase. 6/6 pass, max err 9.5e-7.
The norm distinction is the trap: LayerNorm32 is applied to x.feats (per-voxel over
channels) while SparseGroupNorm32 reshapes [N_b,C] -> [1,C,N_b] per batch item, so its
statistics span channels-in-group AND voxels. Both produce identical shapes, so a mixup
is numerically silent - there is an explicit test asserting the group norm does NOT
match a per-voxel group_norm.
Architecture confirmed from the converted weights rather than constructor defaults:
no rope, no qk_rms_norm, transformer norms non-affine, ResBlock norm1 affine/norm2 not.
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.
