SparseStructureFlowModel verified against upstream at correlation 1.00000000

700/700 params, max abs diff 1.2e-5 on the real 1.3B checkpoint.

Key realisation: the flow models have no sparse conv, so upstream RUNS on CPU torch
with flash-attn swapped for SDPA. That gives a real numerical oracle - unavailable for
the sparse path, where spconv cannot be installed at all.

It was needed. Three bugs survived a loader reporting a perfect 700/700 with zero
missing and zero unmapped keys:
- a parameterless final LayerNorm (no params -> no checkpoint trace) that the output
  was 200x too large without
- rope_phases being complex64, so the rotation is a complex multiply
- qk_rms_norm belonging before rope rather than after

Weight-key matching is necessary but nowhere near sufficient for a port.

CLAUDE.md

@@ -37,3 +37,17 @@ MLX port of TencentARC Pixal3D, on top of the shared `trellis_sparse_mlx` core.
   classify_5d() distinguishes them correctly on the real weights.
 - `rope_phases` ships as a stored tensor — RoPE phases are precomputed, not derived.
 - Converter preserves dtype (fp16 stays fp16); upcasting doubled 24GB for nothing.
+
+## Numerical oracle (important)
+The flow models have NO sparse conv, so upstream RUNS ON CPU TORCH here. Patch
+`pixal3d.modules.attention.modules.scaled_dot_product_attention` with a torch SDPA
+wrapper (permute to [B,H,N,D] and back) and it loads the real checkpoint. Use this to
+diff any flow-model change — do not reason about correctness, measure it.
+
+## Bugs found that weight-matching could NOT catch (loader said 700/700, 0 missing)
+1. Parameterless final F.layer_norm between last block and out_layer. No params -> no
+   checkpoint trace. Without it output is ~200x too large.
+2. rope_phases is complex64 (torch.polar). Rotation is a complex multiply; cos/sin are
+   the real/imag parts, NOT cos(phase).
+3. qk_rms_norm goes BEFORE rope, not after. Non-commutative. One block correlates 0.9998
+   when reversed; compounds to 0.84 over 30 blocks.

README.md

@@ -71,7 +71,36 @@ Weights: 24.04 GB across 19 files (1.3B DiTs at 512/1024 + shape/tex decoders).
 - [x] `upsample` (masked) + `downsample(mode=)` landed in the shared core
 - [x] **Weight converter** — decoders remap KRSC→`[K³,in,out]`, dtype preserved,
       remap verified a pure permutation. Flow models pass through untouched.
-- [ ] DiT blocks: RoPE (stored phases), qk_rms_norm, AdaLN modulation
+- [x] DiT blocks: RoPE (stored complex phases), qk_rms_norm, AdaLN modulation, proj conditioning
+- [x] **`SparseStructureFlowModel` verified against upstream: correlation 1.00000000**,
+      max abs diff 1.2e-5, on the real 1.3B checkpoint (700/700 params)
 - [ ] `SparseConvNeXtBlock3d`, `SparseResBlockC2S3d`, `SparseSpatial2Channel`
 - [ ] Model graphs
 - [ ] End-to-end
+
+
+## Numerical verification
+
+The flow models contain no sparse convolution, which means **upstream runs on CPU torch
+here** — swap flash-attn for `F.scaled_dot_product_attention` and it loads the real
+checkpoint and runs. So unlike the sparse path (where spconv is uninstallable and the
+oracle had to be hand-written), these are diffed against upstream directly:
+
+```
+MLX      : mean +0.18490  std 0.86841
+UPSTREAM : mean +0.18490  std 0.86841
+max abs diff 1.216e-05   correlation 1.00000000
+```
+
+Getting there required finding three bugs that **weight-key matching could not catch** —
+the loader reported a perfect 700/700 with 0 missing and 0 unmapped through all of them:
+
+1. **A parameterless final `LayerNorm`** between the last block and `out_layer`. It has
+   no weights, so it leaves no trace in the checkpoint. Without it the output was ~200x
+   too large (std 187 vs 0.87).
+2. **`rope_phases` is complex64**, built with `torch.polar`. The rotation is a complex
+   multiply and cos/sin are the phase's real/imaginary parts — taking `cos()` of a
+   complex phase is meaningless.
+3. **qk RMS norm is applied BEFORE RoPE, not after.** They do not commute. Reversed, a
+   single block still correlated 0.9998 with upstream; over 30 blocks that compounds to
+   0.84. This one is invisible without an oracle.

pixal3d_mlx/ss_flow.py

@@ -0,0 +1,203 @@
+"""Pixal3D SparseStructureFlowModel in MLX.
+
+The smallest of Pixal3D's four flow checkpoints (5.0 GB) and the one that exercises
+every DiT feature the others use — RoPE, per-head q/k RMS norm, shared AdaLN modulation,
+and `image_attn_mode="proj"` conditioning — so getting this one loading and running
+validates the shared DiT core against real trained weights.
+
+Despite the name it is fully DENSE: it operates on a 16³ voxel grid flattened to 4096
+tokens. No sparse convolution anywhere (the checkpoint contains zero rank-5 tensors).
+
+Config comes from the sibling JSON, not from constructor defaults:
+    resolution 16, in/out 8ch, model 1536, cond 1024, 30 blocks, 12 heads,
+    mlp_ratio 5.3334, pe_mode rope, share_mod, qk_rms_norm (+cross), image_attn_mode proj
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from trellis_sparse_mlx.dit import ModulatedTransformerCrossBlock, TimestepEmbedder
+
+
+class SparseStructureFlowModel(nn.Module):
+    def __init__(
+        self,
+        resolution: int = 16,
+        in_channels: int = 8,
+        out_channels: int = 8,
+        model_channels: int = 1536,
+        cond_channels: int = 1024,
+        num_blocks: int = 30,
+        num_heads: int = 12,
+        mlp_ratio: float = 5.3334,
+        share_mod: bool = True,
+        qk_rms_norm: bool = True,
+        qk_rms_norm_cross: bool = True,
+        image_attn_mode: str = "proj",
+        proj_in_channels: Optional[int] = None,
+        pe_mode: str = "rope",
+        **_ignored,
+    ):
+        super().__init__()
+        self.resolution = resolution
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.model_channels = model_channels
+        self.share_mod = share_mod
+        self.pe_mode = pe_mode
+
+        self.t_embedder = TimestepEmbedder(model_channels)
+        if share_mod:
+            # upstream is Sequential(SiLU, Linear) -> checkpoint key is adaLN_modulation.1
+            self.adaLN_modulation = nn.Linear(model_channels, 6 * model_channels)
+        self.input_layer = nn.Linear(in_channels, model_channels)
+        self.blocks = [
+            ModulatedTransformerCrossBlock(
+                model_channels,
+                cond_channels,
+                num_heads=num_heads,
+                mlp_ratio=mlp_ratio,
+                share_mod=share_mod,
+                use_rope=(pe_mode == "rope"),
+                qk_rms_norm=qk_rms_norm,
+                qk_rms_norm_cross=qk_rms_norm_cross,
+                image_attn_mode=image_attn_mode,
+                proj_in_channels=proj_in_channels,
+            )
+            for _ in range(num_blocks)
+        ]
+        self.out_layer = nn.Linear(model_channels, out_channels)
+        self.rope_phases: Optional[mx.array] = None  # loaded from the checkpoint
+
+    def __call__(self, x: mx.array, t: mx.array, cond) -> mx.array:
+        """x: [B, C, D, H, W]; t: [B]; cond: dict/tuple of (global, proj)."""
+        b, c = x.shape[0], x.shape[1]
+        if c != self.in_channels or list(x.shape[2:]) != [self.resolution] * 3:
+            raise ValueError(
+                f"expected [B,{self.in_channels},{self.resolution}^3], got {x.shape}"
+            )
+        h = x.reshape(b, c, -1).transpose(0, 2, 1)  # [B, N, C]
+        h = self.input_layer(h)
+
+        t_emb = self.t_embedder(t)
+        if self.share_mod:
+            t_emb = self.adaLN_modulation(nn.silu(t_emb))
+
+        phases = self.rope_phases
+        for blk in self.blocks:
+            h = blk(h, t_emb, cond, phases=phases)
+
+        # Parameterless final LayerNorm — `F.layer_norm(h, h.shape[-1:])` upstream.
+        # Easy to miss and impossible to catch by weight-key matching: it has no
+        # parameters, so a loader can report a perfect 700/700 with 0 missing and 0
+        # unmapped while the model is still wrong. The residual stream leaves the last
+        # block at std ~230; without this the output is ~200x too large.
+        mu = mx.mean(h, axis=-1, keepdims=True)
+        var = mx.var(h, axis=-1, keepdims=True)
+        h = (h - mu) * mx.rsqrt(var + 1e-5)
+
+        h = self.out_layer(h)
+        return h.transpose(0, 2, 1).reshape(b, self.out_channels, *[self.resolution] * 3)
+
+
+# ------------------------------------------------------------------ loading
+
+def _remap(k: str) -> str:
+    """Checkpoint key -> module path.
+
+    Upstream wraps a couple of things in nn.Sequential, so its keys carry numeric indices
+    where this implementation uses named attributes:
+        t_embedder.mlp.{0,2}    -> t_embedder.mlp_{0,2}      (Linear, SiLU, Linear)
+        blocks.N.mlp.mlp.{0,2}  -> blocks.N.mlp_{0,2}
+        adaLN_modulation.1      -> adaLN_modulation          (index 0 is the SiLU)
+    """
+    k = k.replace("t_embedder.mlp.0.", "t_embedder.mlp_0.")
+    k = k.replace("t_embedder.mlp.2.", "t_embedder.mlp_2.")
+    k = k.replace(".mlp.mlp.0.", ".mlp_0.")
+    k = k.replace(".mlp.mlp.2.", ".mlp_2.")
+    k = k.replace("adaLN_modulation.1.", "adaLN_modulation.")
+    return k
+
+
+def load(weights_path: str | Path, config_path: str | Path | None = None):
+    """Build from the sibling JSON config and load weights. Returns (model, report)."""
+    wp = Path(weights_path)
+    cp = Path(config_path) if config_path else wp.with_suffix(".json")
+    cfg = json.loads(cp.read_text())
+    args = dict(cfg.get("args", {}))
+    args.pop("dtype", None)
+    args.pop("initialization", None)
+
+    model = SparseStructureFlowModel(**args)
+    w = mx.load(str(wp))
+
+    flat = dict(_flatten(model.parameters()))
+    mapped, unmapped = {}, []
+    for k, v in w.items():
+        if k == "rope_phases":
+            model.rope_phases = v
+            continue
+        m = _remap(k)
+        if m in flat:
+            if flat[m].shape != v.shape:
+                raise ValueError(f"shape mismatch {k} -> {m}: {flat[m].shape} vs {v.shape}")
+            mapped[m] = v
+        else:
+            unmapped.append(f"{k} -> {m}")
+    missing = [k for k in flat if k not in mapped]
+    if mapped:
+        model.update(_unflatten(mapped))
+    return model, {
+        "config": cfg.get("name"),
+        "loaded": len(mapped),
+        "params": len(flat),
+        "missing": missing,
+        "unmapped": unmapped,
+        "rope_phases": None if model.rope_phases is None else model.rope_phases.shape,
+    }
+
+
+def _flatten(tree, prefix=""):
+    if isinstance(tree, dict):
+        for k, v in tree.items():
+            yield from _flatten(v, f"{prefix}{k}.")
+    elif isinstance(tree, list):
+        for i, v in enumerate(tree):
+            yield from _flatten(v, f"{prefix}{i}.")
+    elif isinstance(tree, mx.array):
+        yield prefix[:-1], tree
+
+
+def _unflatten(flat: dict):
+    root: dict = {}
+    for key, val in flat.items():
+        parts = key.split(".")
+        node = root
+        for i, p in enumerate(parts[:-1]):
+            nxt = parts[i + 1]
+            default = [] if nxt.isdigit() else {}
+            if isinstance(node, list):
+                idx = int(p)
+                while len(node) <= idx:
+                    node.append({})
+                if isinstance(default, list) and not isinstance(node[idx], list):
+                    node[idx] = default
+                node = node[idx]
+            else:
+                if p not in node or not isinstance(node[p], (dict, list)):
+                    node[p] = default
+                node = node[p]
+        if isinstance(node, list):
+            idx = int(parts[-1])
+            while len(node) <= idx:
+                node.append(None)
+            node[idx] = val
+        else:
+            node[parts[-1]] = val
+    return root