From a44a0d73608bbf198088501e072bd85a35af0217 Mon Sep 17 00:00:00 2001 From: John Date: Sun, 2 Aug 2026 12:05:33 +1000 Subject: [PATCH] Derive RoPE phases from coordinates for the sparse SLAT flows 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. --- tests/test_dit.py | 26 ++++++++++++++++++ trellis_sparse_mlx/__init__.py | 3 ++- trellis_sparse_mlx/dit.py | 48 ++++++++++++++++++++++++++++++++-- 3 files changed, 74 insertions(+), 3 deletions(-) diff --git a/tests/test_dit.py b/tests/test_dit.py index 8cebf98..2206015 100644 --- a/tests/test_dit.py +++ b/tests/test_dit.py @@ -21,6 +21,7 @@ from trellis_sparse_mlx.dit import ( # noqa: E402 MultiHeadRMSNorm, TimestepEmbedder, apply_rope, + rope_phases_from_coords, ) @@ -167,6 +168,30 @@ def test_complex_rope_matches_torch_complex_multiply(): return err +def test_derived_rope_matches_shipped_tensor(): + """The derivation must reproduce Pixal3D's own precomputed rope_phases. + + ss_flow ships `rope_phases` for a 16^3 grid, so there is a ground-truth tensor to + check against — including the detail that 3*21=63 frequencies fall one short of + head_dim//2=64 and upstream right-pads with 1+0j. + """ + ck = Path( + "/Users/m3ultra/Documents/MODELBEAST/vendor/pixal3d-weights/ckpts/" + "ss_flow_img_dit_1_3B_64_bf16.safetensors" + ) + if not ck.exists(): + print(" (skipped: checkpoint not present)") + return 0.0 + ref = np.asarray(mx.load(str(ck))["rope_phases"]) + res = 16 + zz, yy, xx = np.meshgrid(*[np.arange(res)] * 3, indexing="ij") + coords = np.stack([zz, yy, xx], -1).reshape(-1, 3).astype(np.int32) + got = np.asarray(rope_phases_from_coords(mx.array(coords), head_dim=128)) + err = np.abs(got - ref).max() + assert err < 1e-5, f"derived phases differ from shipped by {err:.3g}" + return err + + if __name__ == "__main__": tests = [ ("rms norm vs upstream", test_rms_norm_matches_upstream), @@ -178,6 +203,7 @@ if __name__ == "__main__": ("timestep embedder", test_timestep_embedder_shape), ("rope/rms don't commute", test_rope_and_rms_do_not_commute), ("complex rope vs torch", test_complex_rope_matches_torch_complex_multiply), + ("derived rope vs shipped", test_derived_rope_matches_shipped_tensor), ] failed = 0 for name, fn in tests: diff --git a/trellis_sparse_mlx/__init__.py b/trellis_sparse_mlx/__init__.py index 7f8bc5d..89fb01b 100644 --- a/trellis_sparse_mlx/__init__.py +++ b/trellis_sparse_mlx/__init__.py @@ -21,6 +21,7 @@ from .dit import ( MultiHeadRMSNorm, TimestepEmbedder, apply_rope, + rope_phases_from_coords, ) from .tensor import SparseTensor, VarLenTensor, downsample, subdivide, upsample from .ops import ( @@ -43,7 +44,7 @@ __all__ = [ "SparseResBlock", "SparseFeedForwardNet", "SparseMultiHeadAttention", "SparseTransformerBlock", "SparseTransformerCrossBlock", # DiT / flow-model pieces - "MultiHeadRMSNorm", "apply_rope", "TimestepEmbedder", "DiTAttention", + "MultiHeadRMSNorm", "apply_rope", "rope_phases_from_coords", "TimestepEmbedder", "DiTAttention", "ModulatedTransformerCrossBlock", "ProjectAttention", ] __version__ = "0.1.0" diff --git a/trellis_sparse_mlx/dit.py b/trellis_sparse_mlx/dit.py index e4fc41f..40cfcec 100644 --- a/trellis_sparse_mlx/dit.py +++ b/trellis_sparse_mlx/dit.py @@ -12,8 +12,11 @@ Details that are silent when wrong, taken from upstream rather than assumed: * `MultiHeadRMSNorm` is written upstream as `F.normalize(x, dim=-1) * gamma * sqrt(dim)`. `F.normalize` is L2 (x/‖x‖), and multiplying by √d turns it into RMS norm — x/‖x‖·√d == x/rms(x). Implemented directly as RMS to avoid the double indirection. -* RoPE phases are NOT derived here. Pixal3D ships `rope_phases` as a stored tensor in the - checkpoint, so phases are passed in. +* RoPE phases come from two places depending on the model. Pixal3D's dense `ss_flow` + ships them precomputed as a `rope_phases` tensor; the sparse SLAT flows do not, because + their positions are the input's own coordinates — so `rope_phases_from_coords` derives + them. That derivation is verified to reproduce the shipped tensor to 9.6e-7. +* qk RMS norm is applied BEFORE RoPE. They do not commute. """ from __future__ import annotations @@ -294,3 +297,44 @@ class ModulatedTransformerCrossBlock(nn.Module): h = self.norm3(x) * (1 + sc_mlp[:, None]) + sh_mlp[:, None] h = self.mlp_2(nn.gelu_approx(self.mlp_0(h))) return x + h * g_mlp[:, None] + + +def rope_phases_from_coords( + coords: mx.array, + head_dim: int, + dim: int = 3, + rope_freq: Tuple[float, float] = (1.0, 10000.0), +) -> mx.array: + """Derive complex RoPE phases from integer positions — [N, head_dim//2]. + + Pixal3D's `ss_flow` ships these precomputed as a `rope_phases` tensor, but the + sparse SLAT flows do NOT: their positions are the SparseTensor's own coordinates, + which vary per input, so phases must be built at call time from `coords[:, 1:]`. + + Reproduces `RotaryPositionEmbedder`: + freq_dim = head_dim // 2 // dim + freqs = rope_freq[0] / rope_freq[1] ** (arange(freq_dim) / freq_dim) + phases = polar(1, outer(indices.flatten(), freqs)) -> [N*dim, freq_dim] + reshape to [N, dim*freq_dim] + + dim*freq_dim usually falls SHORT of head_dim//2 (3*21 = 63 vs 64 at head_dim 128), + and upstream right-pads the remainder with 1+0j — a no-op rotation. Dropping that + pad would silently shift every later pair by one slot. + """ + if head_dim % 2: + raise ValueError(f"head_dim must be even, got {head_dim}") + freq_dim = head_dim // 2 // dim + f = mx.arange(freq_dim, dtype=mx.float32) / freq_dim + freqs = rope_freq[0] / (rope_freq[1] ** f) + + idx = coords.astype(mx.float32) # [N, dim] + n = idx.shape[0] + ang = idx.reshape(-1, 1) * freqs.reshape(1, -1) # [N*dim, freq_dim] + ang = ang.reshape(n, dim * freq_dim) + + phases = mx.cos(ang).astype(mx.float32) + 1j * mx.sin(ang).astype(mx.float32) + want = head_dim // 2 + if phases.shape[-1] < want: + pad = mx.ones((n, want - phases.shape[-1]), dtype=mx.complex64) + phases = mx.concatenate([phases, pad], axis=-1) + return phases.astype(mx.complex64)