# pixal3d_mrp_mlx — working notes

MLX port of TencentARC Pixal3D, on top of the shared `trellis_sparse_mlx` core.

## Measured facts (don't re-derive)
- Pixal3D inherits TRELLIS.2's sparse module, same as LATO.2. `conv_spconv.py` is the
  same code: `stride==1 and padding is None -> spconv.SubMConv3d`.
- EVERY SparseConv3d in the model is `(channels, out, 3)` -> submanifold only. The
  shared core's SubMConv3d covers all of them.
- attn_mode is 'full' everywhere -> SDPA. No flash-attn/xformers needed.
- Container differs from LATO.2: `SparseTensor(VarLenTensor)`. VarLenTensor is feats +
  explicit slice-per-batch layout, no coords. Both are in the shared core now.
- Cache API spelling is `get_spatial_cache`/`register_spatial_cache` (LATO.2 uses
  cache_get/cache_put). Shared core exposes both.
- `_scale` uses `Fraction`, not int, upstream. Shared core keeps scale as an opaque
  tuple so either works.
- Only sparse-conv-using model file is `models/sc_vaes/sparse_unet_vae.py`.

## Remaining ops to port
- `SparseUpsample(2)` — pixal3d/modules/sparse/spatial/basic.py. Cache-paired: needs the
  subdivision stored by its matching SparseDownsample, or an explicit subdivision tensor.
- `SparseSpatial2Channel(2)` — pixal3d/modules/sparse/spatial/spatial2channel.py.
  Sparse pixel-shuffle; also cache-backed.

## Conventions
- Depends on ../trellis_sparse_mlx (editable install). Do NOT vendor a second copy of
  the sparse core.
- `upstream/Pixal3D` is vendored read-only reference; never edit it.
- Weights live at MODELBEAST/vendor/pixal3d-weights (24GB, outside this repo).

## Weights (measured 2026-08-02)
- Every checkpoint ships a sibling .json with {name, args} — the authoritative config.
  Do NOT infer architecture from shapes here (that was needed for LATO.2, not this).
- ss_flow + slat_flow x3 (~20GB, the bulk): ZERO 5-D tensors. Pure transformers, no
  sparse conv. They pass through the converter untouched.
- shape_dec / tex_dec: 40 KRSC sparse kernels each. ss_dec: 20 dense Conv3d.
  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.