8faf6d592a
slat_flow joins ss_flow: max abs diff 9.3e-6, 700/700 params, against upstream running on CPU torch. All three SLAT checkpoints load clean and run (img2shape 512/1024 and imgshape2tex, the last taking 64 in-channels since shape is concatenated). The SLAT flows differ from ss_flow in two ways, both handled in the shared core: tokens are a SparseTensor's voxels so attention runs per batch item, and RoPE phases are NOT shipped - positions are the input's own coordinates, so they are derived at call time by rope_phases_from_coords. tests/oracle_slat.py keeps the CPU-torch oracle harness: it patches both the dense and sparse flash-attn kernels with SDPA equivalents. Use it rather than reasoning about correctness - it has already caught three bugs a perfect 700/700 key match did not.
109 lines
3.6 KiB
Python
109 lines
3.6 KiB
Python
"""Run upstream Pixal3D SLatFlowModel on CPU torch to produce a reference output.
|
|
|
|
The SLAT flows contain no sparse convolution, so upstream is runnable here once the two
|
|
flash-attn kernels are swapped for torch SDPA. That gives a real numerical oracle.
|
|
"""
|
|
import json
|
|
import sys
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
sys.path.insert(0, "/Users/m3ultra/Documents/pixal3d_mrp_mlx/upstream/Pixal3D")
|
|
|
|
import pixal3d.modules.sparse as sp # noqa: E402
|
|
from pixal3d.modules.sparse.basic import VarLenTensor # noqa: E402
|
|
import pixal3d.modules.attention.modules as AM # noqa: E402
|
|
import pixal3d.modules.sparse.attention.modules as SAM # noqa: E402
|
|
|
|
|
|
def _dense_sdpa(q, k, v, **kw):
|
|
return F.scaled_dot_product_attention(
|
|
q.permute(0, 2, 1, 3), k.permute(0, 2, 1, 3), v.permute(0, 2, 1, 3)
|
|
).permute(0, 2, 1, 3)
|
|
|
|
|
|
AM.scaled_dot_product_attention = _dense_sdpa
|
|
|
|
|
|
def _flatten_kv(t):
|
|
"""-> (feats [T,H,C], list of (start,stop)). Handles varlen and dense [N,L,H,C]."""
|
|
if isinstance(t, VarLenTensor):
|
|
return t.feats, [(s.start, s.stop) for s in t.layout]
|
|
if t.dim() == 4: # [N, L, H, C] dense -> flatten, one segment per batch item
|
|
n, l = t.shape[0], t.shape[1]
|
|
return t.reshape(n * l, *t.shape[2:]), [(i * l, (i + 1) * l) for i in range(n)]
|
|
return t, [(0, t.shape[0])]
|
|
|
|
|
|
def _sparse_sdpa(*args):
|
|
if len(args) == 1:
|
|
ref = args[0]
|
|
f = ref.feats
|
|
q, k, v = f[:, 0], f[:, 1], f[:, 2]
|
|
lq = lkv = [(s.start, s.stop) for s in ref.layout]
|
|
elif len(args) == 2:
|
|
ref, kv = args
|
|
q, lq = _flatten_kv(ref)
|
|
kvf, lkv = _flatten_kv(kv)
|
|
k, v = kvf[:, 0], kvf[:, 1]
|
|
else:
|
|
ref, k_, v_ = args
|
|
q, lq = _flatten_kv(ref)
|
|
k, lkv = _flatten_kv(k_)
|
|
v, _ = _flatten_kv(v_)
|
|
outs = []
|
|
for (a, b), (c, d) in zip(lq, lkv):
|
|
o = F.scaled_dot_product_attention(
|
|
q[a:b].permute(1, 0, 2)[None],
|
|
k[c:d].permute(1, 0, 2)[None],
|
|
v[c:d].permute(1, 0, 2)[None],
|
|
)
|
|
outs.append(o[0].permute(1, 0, 2))
|
|
out = torch.cat(outs, 0)
|
|
return ref.replace(out) if isinstance(ref, VarLenTensor) else out
|
|
|
|
|
|
SAM.sparse_scaled_dot_product_attention = _sparse_sdpa
|
|
|
|
from pixal3d.models.structured_latent_flow import SLatFlowModel # noqa: E402
|
|
from safetensors.torch import load_file # noqa: E402
|
|
|
|
CK = (
|
|
"/Users/m3ultra/Documents/MODELBEAST/vendor/pixal3d-weights/ckpts/"
|
|
"slat_flow_img2shape_dit_1_3B_512_bf16"
|
|
)
|
|
cfg = json.load(open(CK + ".json"))
|
|
a = dict(cfg["args"])
|
|
a["dtype"] = "float32"
|
|
m = SLatFlowModel(**a).eval()
|
|
m.load_state_dict(
|
|
{k: v.float() for k, v in load_file(CK + ".safetensors").items()}, strict=False
|
|
)
|
|
|
|
rng = np.random.default_rng(0)
|
|
N = 64
|
|
co = np.concatenate(
|
|
[np.zeros((N, 1), dtype=np.int32), rng.integers(0, 32, (N, 3)).astype(np.int32)], 1
|
|
)
|
|
co = co[np.lexsort((co[:, 3], co[:, 2], co[:, 1], co[:, 0]))]
|
|
fe = rng.standard_normal((N, 32)).astype(np.float32)
|
|
g = (rng.standard_normal((1, 37, 1024)) * 0.1).astype(np.float32)
|
|
pr = (rng.standard_normal((N, 2048)) * 0.1).astype(np.float32)
|
|
t = torch.tensor(np.array([0.5], dtype=np.float32))
|
|
|
|
with torch.no_grad():
|
|
out = m(
|
|
sp.SparseTensor(torch.tensor(fe), torch.tensor(co)),
|
|
t,
|
|
{
|
|
"global": torch.tensor(g),
|
|
"proj": sp.SparseTensor(torch.tensor(pr), torch.tensor(co)),
|
|
},
|
|
)
|
|
o = out.feats.numpy()
|
|
print(f" UPSTREAM slat: {o.shape} mean {o.mean():+.5f} std {o.std():.5f}")
|
|
for n, v in [("ref", o), ("co", co), ("fe", fe), ("g", g), ("pr", pr)]:
|
|
np.save(f"/tmp/px_slat_{n}.npy", v)
|