LATO.2 hardcodes reduce='amax' but Pixal3D's version of the same op defaults to
'mean'. Both now share one implementation, so relying on its default would silently
change every downsampled feature if that default ever moved.
Pixal3D is built on the same TRELLIS.2 sparse module, so the core belongs in one
tested package rather than a copy per port. Verified behaviour-preserving: the
encoder still loads 102/102 params and produces a bit-identical latent
(mean +0.0508, std 0.9186) after the extraction.
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.
Final ordering at 128^3/128ch now tracks GPU core count, which is the evidence the
op is genuinely compute-bound rather than dispatch-bound:
m3ultra (80c) 13.8ms | m2max (38c) 27.3ms | m1ultra (64c) 39.7ms | m1max (32c) 41.4ms
m4pro reachable by ICMP but not ssh (key/user), and it is HANDS OFF per fleet rules,
so it is excluded rather than provisioned.
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.
