DiT blocks: RoPE, per-head RMS norm, AdaLN modulation (23 tests total)

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.

tests/test_dit.py

@@ -0,0 +1,155 @@
+"""DiT block tests against torch.
+
+These layers all have torch counterparts, so unlike the submanifold conv they are checked
+against upstream's real semantics. The upstream forward bodies are reproduced verbatim
+from pixal3d/modules/{attention/modules.py,transformer/modulated.py}.
+"""
+
+import sys
+from pathlib import Path
+
+import mlx.core as mx
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+
+from trellis_sparse_mlx.dit import (  # noqa: E402
+    DiTAttention,
+    ModulatedTransformerCrossBlock,
+    MultiHeadRMSNorm,
+    TimestepEmbedder,
+    apply_rope,
+)
+
+
+def test_rms_norm_matches_upstream(dim=32, heads=4, n=17):
+    """Upstream writes it as F.normalize(x)*gamma*sqrt(dim); we implement RMS directly."""
+    rng = np.random.default_rng(0)
+    x = rng.standard_normal((n, heads, dim)).astype(np.float32)
+    g = rng.standard_normal((heads, dim)).astype(np.float32)
+
+    m = MultiHeadRMSNorm(dim, heads)
+    m.gamma = mx.array(g)
+    got = np.asarray(m(mx.array(x)))
+
+    tx = torch.tensor(x)
+    want = (F.normalize(tx.float(), dim=-1) * torch.tensor(g) * dim**0.5).numpy()
+
+    err = np.abs(got - want).max()
+    assert err < 2e-4, f"rms norm err {err:.3g}"
+    return err
+
+
+def test_rms_norm_is_scale_invariant():
+    """RMS norm must remove input magnitude — catches a plain scale-by-gamma stub."""
+    m = MultiHeadRMSNorm(8, 2)
+    x = np.random.default_rng(1).standard_normal((5, 2, 8)).astype(np.float32)
+    a = np.asarray(m(mx.array(x)))
+    b = np.asarray(m(mx.array(x * 37.0)))
+    err = np.abs(a - b).max()
+    assert err < 1e-3, f"not scale-invariant: {err:.3g}"
+    return err
+
+
+def test_rope_is_a_rotation():
+    """RoPE must preserve the norm of each rotary pair."""
+    rng = np.random.default_rng(2)
+    q = rng.standard_normal((2, 6, 4, 16)).astype(np.float32)
+    ph = rng.uniform(0, 2 * np.pi, (2, 6, 4, 8)).astype(np.float32)
+    rq, _ = apply_rope(mx.array(q), mx.array(q), mx.array(ph))
+    rq = np.asarray(rq)
+    n0 = np.sqrt(q[..., 0::2] ** 2 + q[..., 1::2] ** 2)
+    n1 = np.sqrt(rq[..., 0::2] ** 2 + rq[..., 1::2] ** 2)
+    err = np.abs(n0 - n1).max()
+    assert err < 1e-4, f"rope changed pair norms by {err:.3g}"
+    # and it must actually rotate
+    assert np.abs(rq - q).max() > 1e-3, "rope was a no-op"
+    return err
+
+
+def test_attention_matches_torch(ch=64, heads=8, n=13, b=2):
+    rng = np.random.default_rng(3)
+    x = rng.standard_normal((b, n, ch)).astype(np.float32)
+    wq = rng.standard_normal((ch * 3, ch)).astype(np.float32) * 0.05
+    bq = rng.standard_normal((ch * 3,)).astype(np.float32) * 0.05
+    wo = rng.standard_normal((ch, ch)).astype(np.float32) * 0.05
+    bo = rng.standard_normal((ch,)).astype(np.float32) * 0.05
+
+    a = DiTAttention(ch, heads)
+    a.to_qkv.weight, a.to_qkv.bias = mx.array(wq), mx.array(bq)
+    a.to_out.weight, a.to_out.bias = mx.array(wo), mx.array(bo)
+    got = np.asarray(a(mx.array(x)))
+
+    d = ch // heads
+    tx = torch.tensor(x)
+    qkv = F.linear(tx, torch.tensor(wq), torch.tensor(bq)).reshape(b, n, 3, heads, d)
+    q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]
+    o = F.scaled_dot_product_attention(
+        q.permute(0, 2, 1, 3), k.permute(0, 2, 1, 3), v.permute(0, 2, 1, 3)
+    )
+    o = o.permute(0, 2, 1, 3).reshape(b, n, ch)
+    want = F.linear(o, torch.tensor(wo), torch.tensor(bo)).detach().numpy()
+
+    err = np.abs(got - want).max()
+    assert err < 2e-4, f"attention err {err:.3g}"
+    return err
+
+
+def test_modulation_gates_actually_gate(ch=32, heads=4, n=7, b=2):
+    """gate=0 must make the block an identity on the modulated paths."""
+    blk = ModulatedTransformerCrossBlock(ch, ch, heads, share_mod=True)
+    x = mx.array(np.random.default_rng(4).standard_normal((b, n, ch)).astype(np.float32))
+    ctx = mx.array(np.zeros((b, 5, ch), dtype=np.float32))
+    # zero everything, then force the cross-attn output to zero via zero out-proj
+    blk.modulation = mx.zeros((6 * ch,))
+    blk.cross_attn.to_out.weight = mx.zeros_like(blk.cross_attn.to_out.weight)
+    blk.cross_attn.to_out.bias = mx.zeros_like(blk.cross_attn.to_out.bias)
+    mod = mx.zeros((b, 6 * ch))
+    out = np.asarray(blk(x, mod, ctx))
+    err = np.abs(out - np.asarray(x)).max()
+    assert err < 1e-4, f"gates did not zero the residual branches: {err:.3g}"
+    return err
+
+
+def test_norm_affine_asymmetry():
+    """norm1/norm3 non-affine, norm2 affine — silent if swapped."""
+    blk = ModulatedTransformerCrossBlock(16, 16, 2, share_mod=True)
+    assert not blk.norm1.affine, "norm1 should be non-affine"
+    assert blk.norm2.affine, "norm2 SHOULD be affine"
+    assert not blk.norm3.affine, "norm3 should be non-affine"
+    return 0.0
+
+
+def test_timestep_embedder_shape():
+    te = TimestepEmbedder(64, 32)
+    out = te(mx.array(np.array([0.0, 0.5, 1.0], dtype=np.float32)))
+    assert out.shape == (3, 64), out.shape
+    assert np.isfinite(np.asarray(out)).all()
+    return 0.0
+
+
+if __name__ == "__main__":
+    tests = [
+        ("rms norm vs upstream", test_rms_norm_matches_upstream),
+        ("rms scale invariance", test_rms_norm_is_scale_invariant),
+        ("rope is a rotation", test_rope_is_a_rotation),
+        ("dit attention vs torch", test_attention_matches_torch),
+        ("modulation gates", test_modulation_gates_actually_gate),
+        ("norm affine asymmetry", test_norm_affine_asymmetry),
+        ("timestep embedder", test_timestep_embedder_shape),
+    ]
+    failed = 0
+    for name, fn in tests:
+        try:
+            err = fn()
+            print(f"  PASS  {name:26s} (max err {err:.2e})")
+        except AssertionError as e:
+            print(f"  FAIL  {name:26s} {e}")
+            failed += 1
+        except Exception as e:  # noqa: BLE001
+            print(f"  ERROR {name:26s} {type(e).__name__}: {e}")
+            failed += 1
+    print(f"\n{len(tests)-failed}/{len(tests)} passed")
+    sys.exit(1 if failed else 0)

trellis_sparse_mlx/__init__.py

@@ -14,6 +14,13 @@ than vendoring its own copy.
 """
 
 from .conv import SubMConv3d, build_indice_map
+from .dit import (
+    DiTAttention,
+    ModulatedTransformerCrossBlock,
+    MultiHeadRMSNorm,
+    TimestepEmbedder,
+    apply_rope,
+)
 from .tensor import SparseTensor, VarLenTensor, downsample, subdivide, upsample
 from .ops import (
     LayerNorm32,
@@ -34,5 +41,8 @@ __all__ = [
     "SparseLinear", "LayerNorm32", "SparseGroupNorm32", "SparseSiLU", "SparseGELU",
     "SparseResBlock", "SparseFeedForwardNet", "SparseMultiHeadAttention",
     "SparseTransformerBlock", "SparseTransformerCrossBlock",
+    # DiT / flow-model pieces
+    "MultiHeadRMSNorm", "apply_rope", "TimestepEmbedder", "DiTAttention",
+    "ModulatedTransformerCrossBlock",
 ]
 __version__ = "0.1.0"

trellis_sparse_mlx/dit.py

@@ -0,0 +1,240 @@
+"""Modulated (DiT-style) transformer pieces used by the flow models in this family.
+
+Pixal3D's four flow checkpoints (~20GB — the bulk of its download) are pure transformers
+with no sparse convolution at all, built from these blocks. LATO.2 has flow models too
+(`vertex_structured_flow`, `topo_flow`), so they live in the shared core rather than in
+either port.
+
+Details that are silent when wrong, taken from upstream rather than assumed:
+
+* In `ModulatedTransformerCrossBlock`, `norm1` and `norm3` are NON-affine but `norm2` IS
+  affine (`elementwise_affine=True`). Same shapes either way.
+* `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.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Optional, Tuple
+
+import mlx.core as mx
+import mlx.nn as nn
+
+
+class MultiHeadRMSNorm(nn.Module):
+    """Per-head RMS norm over head_dim, with a [heads, dim] gain."""
+
+    def __init__(self, dim: int, heads: int):
+        super().__init__()
+        self.gamma = mx.ones((heads, dim))
+        self.eps = 1e-12
+
+    def __call__(self, x: mx.array) -> mx.array:
+        # x: [..., heads, dim]
+        dt = x.dtype
+        f = x.astype(mx.float32)
+        rms = mx.sqrt(mx.mean(f * f, axis=-1, keepdims=True) + self.eps)
+        return ((f / rms) * self.gamma).astype(dt)
+
+
+def apply_rope(q: mx.array, k: mx.array, phases: mx.array) -> Tuple[mx.array, mx.array]:
+    """Rotate q/k by precomputed `phases`.
+
+    `phases` is [..., head_dim/2] (or broadcastable) giving the angle per rotary pair.
+    Pairs are (even, odd) along the last axis, matching upstream's interleaved layout.
+    """
+    cos, sin = mx.cos(phases), mx.sin(phases)
+
+    def rot(t: mx.array) -> mx.array:
+        dt = t.dtype
+        f = t.astype(mx.float32)
+        a, b = f[..., 0::2], f[..., 1::2]
+        ra = a * cos - b * sin
+        rb = a * sin + b * cos
+        out = mx.stack([ra, rb], axis=-1).reshape(f.shape)
+        return out.astype(dt)
+
+    return rot(q), rot(k)
+
+
+class TimestepEmbedder(nn.Module):
+    """Sinusoidal timestep embedding -> 2-layer MLP, as in DiT."""
+
+    def __init__(self, hidden_size: int, frequency_embedding_size: int = 256):
+        super().__init__()
+        self.frequency_embedding_size = frequency_embedding_size
+        self.mlp_0 = nn.Linear(frequency_embedding_size, hidden_size)
+        self.mlp_2 = nn.Linear(hidden_size, hidden_size)
+
+    @staticmethod
+    def timestep_embedding(t: mx.array, dim: int, max_period: int = 10000) -> mx.array:
+        half = dim // 2
+        freqs = mx.exp(
+            -math.log(max_period) * mx.arange(half, dtype=mx.float32) / half
+        )
+        args = t.astype(mx.float32)[:, None] * freqs[None]
+        emb = mx.concatenate([mx.cos(args), mx.sin(args)], axis=-1)
+        if dim % 2:
+            emb = mx.concatenate([emb, mx.zeros((emb.shape[0], 1))], axis=-1)
+        return emb
+
+    def __call__(self, t: mx.array) -> mx.array:
+        e = self.timestep_embedding(t, self.frequency_embedding_size)
+        return self.mlp_2(nn.silu(self.mlp_0(e)))
+
+
+class _LN(nn.Module):
+    """LayerNorm computed in fp32, optional affine — upstream's LayerNorm32."""
+
+    def __init__(self, dim: int, affine: bool, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.affine = affine
+        if affine:
+            self.weight = mx.ones((dim,))
+            self.bias = mx.zeros((dim,))
+
+    def __call__(self, x: mx.array) -> mx.array:
+        dt = x.dtype
+        f = x.astype(mx.float32)
+        f = (f - mx.mean(f, -1, keepdims=True)) * mx.rsqrt(
+            mx.var(f, -1, keepdims=True) + self.eps
+        )
+        if self.affine:
+            f = f * self.weight + self.bias
+        return f.astype(dt)
+
+
+class DiTAttention(nn.Module):
+    """Self or cross attention with optional RoPE and per-head q/k RMS norm."""
+
+    def __init__(
+        self,
+        channels: int,
+        num_heads: int,
+        ctx_channels: Optional[int] = None,
+        attn_type: str = "self",
+        qkv_bias: bool = True,
+        use_rope: bool = False,
+        qk_rms_norm: bool = False,
+    ):
+        super().__init__()
+        self.channels, self.num_heads = channels, num_heads
+        self.head_dim = channels // num_heads
+        self.scale = self.head_dim**-0.5
+        self._type = attn_type
+        self.use_rope = use_rope
+        self.qk_rms_norm = qk_rms_norm
+        ctx = ctx_channels if ctx_channels is not None else channels
+        if attn_type == "self":
+            self.to_qkv = nn.Linear(channels, channels * 3, bias=qkv_bias)
+        else:
+            self.to_q = nn.Linear(channels, channels, bias=qkv_bias)
+            self.to_kv = nn.Linear(ctx, channels * 2, bias=qkv_bias)
+        if qk_rms_norm:
+            self.q_rms_norm = MultiHeadRMSNorm(self.head_dim, num_heads)
+            self.k_rms_norm = MultiHeadRMSNorm(self.head_dim, num_heads)
+        self.to_out = nn.Linear(channels, channels)
+
+    def __call__(
+        self,
+        x: mx.array,
+        context: Optional[mx.array] = None,
+        phases: Optional[mx.array] = None,
+    ) -> mx.array:
+        b, n, _ = x.shape
+        h, d = self.num_heads, self.head_dim
+        if self._type == "self":
+            qkv = self.to_qkv(x).reshape(b, n, 3, h, d)
+            q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]
+        else:
+            if context is None:
+                raise ValueError("cross attention needs a context")
+            m = context.shape[1]
+            q = self.to_q(x).reshape(b, n, h, d)
+            kv = self.to_kv(context).reshape(b, m, 2, h, d)
+            k, v = kv[:, :, 0], kv[:, :, 1]
+        if self.use_rope and phases is not None:
+            q, k = apply_rope(q, k, phases)
+        if self.qk_rms_norm:
+            q, k = self.q_rms_norm(q), self.k_rms_norm(k)
+        q = q.transpose(0, 2, 1, 3)
+        k = k.transpose(0, 2, 1, 3)
+        v = v.transpose(0, 2, 1, 3)
+        o = mx.fast.scaled_dot_product_attention(q, k, v, scale=self.scale)
+        return self.to_out(o.transpose(0, 2, 1, 3).reshape(b, n, self.channels))
+
+
+class ModulatedTransformerCrossBlock(nn.Module):
+    """AdaLN-modulated self-attn -> cross-attn -> FFN.
+
+    Modulation is six chunks (shift/scale/gate for MSA and MLP). With `share_mod` the
+    block holds its own `modulation` parameter that is ADDED to the incoming `mod`;
+    otherwise it derives them via its own `adaLN_modulation`.
+
+    Only the self-attention and MLP are modulated — the cross-attention path is not,
+    which is why `norm2` is the affine one.
+    """
+
+    def __init__(
+        self,
+        channels: int,
+        ctx_channels: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        share_mod: bool = False,
+        use_rope: bool = False,
+        qk_rms_norm: bool = False,
+        qk_rms_norm_cross: bool = False,
+    ):
+        super().__init__()
+        self.share_mod = share_mod
+        self.norm1 = _LN(channels, affine=False, eps=1e-6)
+        self.norm2 = _LN(channels, affine=True, eps=1e-6)
+        self.norm3 = _LN(channels, affine=False, eps=1e-6)
+        self.self_attn = DiTAttention(
+            channels, num_heads, use_rope=use_rope, qk_rms_norm=qk_rms_norm
+        )
+        self.cross_attn = DiTAttention(
+            channels,
+            num_heads,
+            ctx_channels=ctx_channels,
+            attn_type="cross",
+            qk_rms_norm=qk_rms_norm_cross,
+        )
+        hidden = int(channels * mlp_ratio)
+        self.mlp_0 = nn.Linear(channels, hidden)
+        self.mlp_2 = nn.Linear(hidden, channels)
+        if share_mod:
+            self.modulation = mx.zeros((6 * channels,))
+        else:
+            self.adaLN_modulation_1 = nn.Linear(channels, 6 * channels)
+
+    def __call__(
+        self,
+        x: mx.array,
+        mod: mx.array,
+        context: mx.array,
+        phases: Optional[mx.array] = None,
+    ) -> mx.array:
+        if self.share_mod:
+            m = (self.modulation + mod).astype(mod.dtype)
+        else:
+            m = self.adaLN_modulation_1(nn.silu(mod))
+        c = m.shape[-1] // 6
+        sh_msa, sc_msa, g_msa, sh_mlp, sc_mlp, g_mlp = (
+            m[..., i * c : (i + 1) * c] for i in range(6)
+        )
+
+        h = self.norm1(x) * (1 + sc_msa[:, None]) + sh_msa[:, None]
+        x = x + self.self_attn(h, phases=phases) * g_msa[:, None]
+
+        x = x + self.cross_attn(self.norm2(x), context)
+
+        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]