Two silent bugs, both found by diffing against upstream running on CPU torch. The flow
models have no sparse conv, so upstream is runnable here with flash-attn swapped for
SDPA - a real numerical oracle, unlike the sparse path.
1. rope_phases ships as COMPLEX64 (torch.polar), so the rotation is a complex multiply
and cos/sin are the phase's real/imag parts. Taking cos() of a complex phase was
completely wrong. Verified against torch's view_as_complex formulation to 1.2e-7.
2. Upstream applies qk RMS norm BEFORE RoPE; I had it reversed. They do not commute -
RMS applies a per-component gain, RoPE rotates within pairs. Reversed, one block
still correlated 0.9998, which compounded to 0.84 across 30 blocks.
After both: SparseStructureFlowModel matches upstream at correlation 1.00000000,
max abs diff 1.2e-5, on the real 1.3B checkpoint.
Unlocks all four Pixal3D flow checkpoints (~20GB) at once - they are pure transformers
with no sparse conv. LATO.2 has flow models too (vertex_structured_flow, topo_flow), so
these belong in the shared core rather than either port.
Three details taken from upstream rather than assumed, each silent when wrong:
- norm1/norm3 are NON-affine but norm2 IS affine in the modulated cross block. There is
an explicit test asserting that asymmetry.
- MultiHeadRMSNorm is written upstream as F.normalize(x)*gamma*sqrt(dim). F.normalize is
L2, and the sqrt(d) turns it into RMS - implemented directly as RMS and verified equal
to the upstream formulation to 9.5e-7.
- RoPE phases are NOT derived: Pixal3D ships rope_phases as a stored tensor, so they are
passed in. Tested that the rotation preserves per-pair norms and is not a no-op.
Also tested: gates at zero make the block an identity on its residual branches.
Closes the op gap for Pixal3D.
downsample now takes mode='max'|'mean'. This differs BETWEEN MODELS and is silent when
wrong: LATO.2 hardcodes amax, Pixal3D parameterises it and defaults to mean. Callers
must be explicit.
upsample is a MASKED expansion, distinct from subdivide: subdivision.feats is an
[N, factor**3] 0/1 mask, so a voxel emits 0-8 children rather than always 8. Pixal3D's
decoder uses it to grow the occupied set selectively. Child-index decoding is
subidx // factor**i % factor - axis 0 fastest-varying, the REVERSE of subdivide's
C-order, which is exactly the kind of mixup that misplaces every child silently.
Tests cover both modes differing, masked expansion landing at the right offsets, and a
voxel with no flagged children contributing nothing.
Extracted from lato.2_mrp_mlx. The same sparse module underlies LATO.2, Pixal3D and
the rest of the TRELLIS.2 family, and all of them are blocked on Apple Silicon by the
same single op - submanifold conv - so it belongs in one tested package rather than
vendored per port.
13/13 tests: 7 for the conv against a hand-written reference (spconv is uninstallable
here so there is no upstream oracle), 6 for the remaining layers against torch.
SubMConv3d runs 13.8ms at 128^3/128ch on m3ultra, 5.9x faster than the obvious
per-offset loop.
