202ddde31d
The plan was to port the fused Metal spconv kernel and the 15 mx.compile sites down
from trellis-2-mrp-mlx into the shared core. Measured on the full textured cascade,
both are dead ends. PROFILE.md has the numbers; the short version:
BY BACKEND (204.3s total, peak 37.1GB)
DiT transformer flows ~173s ~85%
sparse conv (decoders) ~26s ~13%
torch/MPS (DINOv3+NAF) 4.1s 2%
o_voxel native 0.8s 0.4%
The Metal kernel can only touch the 13%. Verified from the checkpoints: the four flow
models contain ZERO 5-D tensors, so they never call sparse conv at all - only shape_dec
(40), tex_dec (40) and ss_dec (20, dense) do. A free 2x on ALL sparse conv returns 6%
of runtime. The kernel is still the right port if peak memory ever matters, which is
its stated prize, but it is not a speed fix.
mx.compile is no better, because the DiT loop is compute-bound rather than
dispatch-bound - the opposite of the launch-latency problem that motivated the fused
gather-matmul in the shared core:
scaling 25% tokens -> 0.15x, 50% -> 0.36x, 100% -> 1.00x (super-linear)
per block attention 1062 GFLOP 58.8% (O(n^2)), MLP 496, proj 248
12.7 TFLOP/s achieved at 142 ms/block
direct test eager 23.9 ms vs mx.compile 23.7 ms -> 0.8%, i.e. nothing
Attention already routes through mx.fast.scaled_dot_product_attention. mx.compile also
wants pure array-in/array-out functions while the sparse path threads SparseTensor
objects with Python-side layout, so it would mean restructuring for a measured ~0%.
The real lever is TOKEN COUNT: attention is O(n^2), half the tokens ran 2.8x faster,
and refine_coords already backs the grid off past max_num_tokens. That knob trades
resolution for time honestly.
m4pro cannot run this: peak 37.1GB against 24GB of RAM, and the weights alone are 24GB
before activations. Capacity, not tuning.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
63 lines
1.1 KiB
JSON
63 lines
1.1 KiB
JSON
{
|
|
"stages": {
|
|
"ss_flow+ss_dec": {
|
|
"seconds": 18.0,
|
|
"peak_gb": 6.7,
|
|
"delta_gb": 6.7
|
|
},
|
|
"cond_512": {
|
|
"seconds": 1.6,
|
|
"peak_gb": 6.7,
|
|
"delta_gb": 0.0
|
|
},
|
|
"LR slat_flow": {
|
|
"seconds": 12.5,
|
|
"peak_gb": 12.4,
|
|
"delta_gb": 5.7
|
|
},
|
|
"refine(dec)": {
|
|
"seconds": 3.4,
|
|
"peak_gb": 15.5,
|
|
"delta_gb": 3.1
|
|
},
|
|
"cond_1024+NAF": {
|
|
"seconds": 0.9,
|
|
"peak_gb": 15.5,
|
|
"delta_gb": 0.0
|
|
},
|
|
"HR slat_flow": {
|
|
"seconds": 90.9,
|
|
"peak_gb": 23.6,
|
|
"delta_gb": 8.1
|
|
},
|
|
"shape_dec": {
|
|
"seconds": 11.6,
|
|
"peak_gb": 28.4,
|
|
"delta_gb": 4.8
|
|
},
|
|
"fdg_to_mesh": {
|
|
"seconds": 0.8,
|
|
"peak_gb": 28.4,
|
|
"delta_gb": 0.0
|
|
},
|
|
"cond_tex+NAF": {
|
|
"seconds": 1.6,
|
|
"peak_gb": 28.4,
|
|
"delta_gb": 0.0
|
|
},
|
|
"tex_flow+tex_dec": {
|
|
"seconds": 63.0,
|
|
"peak_gb": 37.1,
|
|
"delta_gb": 8.7
|
|
}
|
|
},
|
|
"by_backend": {
|
|
"dit": 121.4,
|
|
"torch": 4.1,
|
|
"sparse_conv": 15.0,
|
|
"native": 0.8,
|
|
"mixed": 63.0
|
|
},
|
|
"total": 204.3,
|
|
"peak_gb": 37.1
|
|
} |