SparseDiTAttention / SparseProjectAttention / ModulatedSparseTransformerCrossBlock -
the sparse counterparts of the dense DiT pieces, attending within each batch item.
The modulation needs care: it is per batch item ([B, 6C]) while features are a flat
[N, C] stack, so each row must pick up its own item's shift/scale/gate. Broadcasting
would silently apply item 0's modulation to everything when B == 1.
Verified end to end: slat_flow matches upstream at correlation 1.00000000.
ss_flow ships rope_phases precomputed; the SLAT flows do not, because their positions
are the input SparseTensor's own coordinates and vary per input. rope_phases_from_coords
reproduces RotaryPositionEmbedder, including the pad detail: 3*21=63 frequencies fall one
short of head_dim//2=64 and upstream right-pads with 1+0j, so dropping it would shift
every later pair by a slot.
Verified against ss_flow's own shipped tensor to 9.6e-7 - a real ground truth, since the
same embedder produced it.
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.
spconv KRSC [O,kz,ky,kx,I] and torch Conv3d [O,I,kz,ky,kx] are both rank 5 and cannot
be separated by rank. Both LATO.2 and Pixal3D contain a mix, and guessing wrong is
silent - the reshape succeeds and the output is quietly garbage. Verified against
Pixal3D's real weights: ss_dec classified conv3d (20), shape_dec/tex_dec krsc (40 each).
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.
Pixal3D refactored the container into a VarLenTensor base (feats + explicit
slice-per-batch layout, no coords) with SparseTensor(VarLenTensor) adding
coordinates. LATO.2 has only the combined class. Modelling the split here means one
package serves both without either port adapting at every call site.
Also exposes Pixal3D's cache spelling (get/register_spatial_cache) alongside LATO.2's
(cache_get/cache_put). 13/13 tests unchanged.
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.
