Repository navigation
feat: add SAM 3D Objects reconstruction and geometry exports - #2578
Merged
Merged
Conversation
Port the two-stage Objects flow, sparse decoders, FlexiCubes extraction, shared DINO conditioning, and optional bundled MoGe depth. Preserve BF16 main weights and F32 depth weights, validate exact checkpoint schemas, and expose typed mask/pointmap sidecars with OBJ and Gaussian PLY outputs. Add authenticated small upstream oracles, full-output capture scaffolding, and strict geometry/input regressions. Eight lightweight repository contracts and independent static implementation/security reviews pass. Native compilation, full reference execution, Rust parity, and real default CLI acceptance remain pending the shared-machine validation schedule. Refs #2381
inureyes
added
status:in-progress
CI compilation found that borrowing a fallible UniquePtr expression before applying the question-mark operator prevented deref coercion to MlxArray at fourteen call sites. Dereference those results explicitly and bind the depth RGB input before calling MoGe. Remove an unused latent-condition re-export. Formatting and whitespace checks pass. CI will recheck compilation; local native work remains paused during the exclusive upstream reference run. Refs #2381
Keep DINO tests in their original module after extracting the adapter, explicitly dereference fallible arrays in numerical tests, and address Clippy's iterator and optional-reference diagnostics. Preserve upstream normalization literals with scoped precision annotations and retain the empty-chunk guard during mesh readback. Formatting and whitespace checks pass. These changes address the current CI diagnostics; full compilation and model parity remain pending. Refs #2381
Record completed unchanged-upstream seed-42 default and RGBA/dropout reference metadata, retaining numerical arrays in ignored storage. Add independently reviewed process/source/binary/device authentication, a bounded Rust capture driver, and strict full-set coordinate, latent, and mesh comparison. Seventeen mock-only harness tests passed. Both upstream references completed with finite full outputs and zero child/wrapper/guard exits. Actual Rust parity and default CLI execution remain pending; these receipts do not claim those gates passed. Refs #2381
Use the upstream fast attention operation so padded tokens are excluded instead of receiving additive zero scores. Keep the regression oracle unchanged; GPU retesting and full model parity remain pending.
Use upstream fused biased projections in the Objects DINO policy and evict recomputable MoGe neighbor maps under the existing memory limits. Add distinguishing projection and cache-equivalence regressions plus ignored stage diagnostics. Authenticate the dense writer’s canonical empty instances sidecar in raw and CLI captures. The corrected capture helpers pass 32 tests; production rebuild and real-model parity remain pending.
Bound convolution, pooling and refinement phases using checked live-buffer estimates. Reuse borrowed neighbor maps, release evaluated scratch buffers and derive a finite supported budget from validated input dimensions while retaining the initial-grid admission checks. Add the conservative published-pyramid planning regression. Register test-only diagnostic variables and include the reviewed upstream stage generator. Environment checks and 15 focused checker tests pass; corrected sparse runtime tests and full-model parity remain pending.
Use Objects-only ratio-first nearest sampling and adaptive separable resize order to match the authenticated upstream crop tensors. Preserve the exact Power primitive and operand order in compiled tanh GELU for flow and decoder feed-forward layers. Add distinguishing authenticated resize and activation oracles, with separate original capture and portable reproduction identities. Other model families retain their existing interpolation and activation defaults. Full model parity remains pending the corrected rebuild.
Reject infeasible source conversions and aspect-derived crop, square, interpolation, and point-map buffers before allocation. Use checked conservative live-byte accounting with a 1 GiB preprocessing guard; preserve the complete published input and all model arithmetic. Share mask bounds between early validation and preprocessing, and account for automatic depth only when the bundle contains a depth model. Add a pure planner regression for the valid published geometry and a hostile 32 Mi-pixel extreme-aspect input.
The released Objects depth quantization scale does not determine its sparse grid depth. Pin the captured row bounds and complete 5,598,663,168-byte phase estimate, admit it within a finite 8 GiB support envelope, and preserve per-input accounting and oversized-input guards. Add an ignored 25-step structure diagnostic with exact timestep representations, embedding boundaries, states, velocities, and raw decoder output to localize the remaining full-run coordinate mismatch. Production sampling is unchanged. Fix the CI-reported equivalent Option spelling. Host provenance and planner arithmetic checks pass; diagnostic source review passes. Updated Rust compilation and runtime parity remain pending. Refs #2381.
Upstream submits the complete five-output structure state in one async evaluation. Match that boundary and its occupancy-logit submission through existing core APIs, preserving the sampling arithmetic and random keys. Replace the perturbing diagnostic with a deferred 183-tensor trajectory capture. It retains handles during sampling, reconstructs time metadata after occupancy, and serializes only after production boundaries have completed. The failed instrumented reference and canceled build remain preserved; actual full-run parity is still pending. Refs #2381.
Prevent one wave from reusing shared reduction storage for variance before every wave has read the mean. The pre-fix tiny probe was stable, so this does not establish the cause of full-model drift. Source ordering was independently reviewed. Native rebuild and runtime validation follow this frozen change. Refs #2381.
Check nonzero-mean normalization and repeated results across four widths and F32, F16, and BF16. Record the overlay ordering fix and the bounded pre-fix probe's nonreproduction without claiming a confirmed cause of model drift. Formatting and diff checks pass; execution follows the rebuilt native library. Refs #2381.
Add an ignored diagnostic for the remaining automatic-depth parity mismatch. Capture the actual depth point map, nonfinite mask, normalization metadata, prepared inputs, point-patch tokens and projections, and final condition without running either flow stage or changing production math. The 17-tensor schema and test-only accessors passed independent source review against the authenticated upstream capture. Runtime comparison follows the frozen test build. Refs #2381.
Preserve the parity tool's startup configuration while satisfying field_reassign_with_default in the required local all-test clippy check. The prior lint failure was present in the issue base; the prompt-cache setting remains unchanged. Formatting and diff checks pass. The warmed clippy retry follows this source freeze. Refs #2381.
Add test-only scoped observations for the encoder, heads, refinements, and geometry solver to locate the default Objects point-map mismatch. Preserve nonfinite values and corroborate the observed result against the uninstrumented capture before interpreting stage differences.
Avoid the deprecated safetensors size method for the restricted F32, I32, and BOOL diagnostic types.
Label the two cheap encoder preprocessing observations separately from tensors captured directly from the production inference.
Select upstream resize ordering, positional interpolation, DINO projection and activation arithmetic, and UV scalar precision only for Objects auto-depth. Add authentic small interpolation oracles and regressions that preserve the standard MoGe and DINO paths. Restore the parity harness's default-based server initialization with a local lint allowance because optional features add private fields that prevent struct-update syntax.
Use the pinned MoGe eager initialization and thirty compiled LM steps only for bundled Objects free-focal recovery. Preserve standalone and fixed-focal behavior, validate the native boundary, and compare every solver state against authenticated CPU and ROCm saved-input oracles.
Check native tensor byte products before creating focal solver intermediates, including diagnostic trace storage. Exercise both safe entry points with lazy broadcast metadata that exceeds representable byte sizes.
Resolve the MLX stack overload explicitly so the focal solver does not select the surrounding CXX bridge function for initializer-list inputs.
inureyes
added
status:review
Record authenticated default and RGBA reconstruction, actual CLI output validation, reproducible reference tooling, and bilingual implementation reports. Update the dense prediction documentation to describe geometry outputs and the checked allocation policies.
inureyes
marked this pull request as ready for review
October 23, 2026 20:39
This was referenced Oct 24, 2026
inureyes
added a commit
that referenced
this pull request
Oct 26, 2026
) ## Summary The bf16 addmm rounding oracles in `dinov2_layers::policy` and `vision::dense` fail on Metal because they were captured on ROCm gfx1151 and x86_64 Linux CPU (#2551 and #2578 were validated on ROCm, not Metal). The failures come from which MLX kernel runs on each backend. This PR keys each oracle on a property the test can detect, without a production change. It also fixes a default-device race that made one failure vary between runs. ## Root cause - **addmm at one row on Metal rounds twice.** MLX's Metal `AddMM::eval_gpu` sends `min(M, N) == 1` to `gemv_axbpy`. The `GEMVKernel` epilogue in `mlx/backend/metal/kernels/gemv.h` casts the F32 accumulator to bf16 before adding the bias, so the result equals a separate add. Two or more rows go to steel GEMM or `gemv_wide` (`gemv_wide_config` rejects `M <= 1` on every architecture generation), and those apply the bias before the single rounding. CPU and ROCm copy C into the output and run one GEMM with beta, so they round once at every row count. - **Upstream MLX gives the same values.** Upstream `mlx` 0.32.4 from PyPI on this M1 Ultra (`applegpu_g13d`): GPU `mx.addmm` at 1 row gives `1.9375` (separate), at 2, 5, 16 and 64 rows it gives `1.9296875` (fused); the CPU stream gives `1.9296875` at every row count. The fc2 halfway case gives `1.0` at `[1,1,4]` and `1.0078125` at `[2,1,4]`, `[1,2,4]` and `[2,4]`. - **The MLX bump in #2614 did not cause it.** `gemv.h` and `AddMM::eval_gpu` are byte-identical between the previous pin `bff10220` and the current pin `7ca40033`. Upstream `main` still has the cast before the bias add, and the 0.32.4 wheel, which predates both pins, reproduces the failure. I found no existing ml-explore/mlx issue for it. - **Sapiens2 compiled RoPE.** The fixture's compiled value `-4.5597553e-6` is a ROCm capture. Upstream `mlx` 0.32.4 on this host gives `-4.589557647705078e-6` for both the compiled and the eager graph, on GPU and on CPU, and the repository generator `tests/fixtures/sapiens2/generate_rope_oracle.py` stops with "compiled and eager RoPE unexpectedly agree". - **SAM 3D Body MHR CPU geometry.** The fixture was captured with MLX CPU on x86_64 Linux (glibc). On macOS, MLX CPU computes atan2, sin and cos through Apple's math library (`mlx/backend/cpu/simd/accelerate_simd.h`). Of the 388 elements in `body_params`, `hand_params`, `fk_state` and `fk_global_rot`, 27 differ by 1 or 2 ULP, and every other field matches bit for bit. Seven of the 27 are single atan2 calls. For those seven, upstream `mx.arctan2` on this host's CPU returns exactly mlxcel's values, and it differs from the fixture at exactly those seven. - **The varying fifth failure is a device race.** The MHR test holds `DefaultDeviceGuard::cpu()`. Tests that read the default device without `lock_default_device()` could run on the CPU in the middle of their span. That is why `compiled_objects_gelu_tanh_matches_authentic_power_oracle` failed only in parallel, and `sapiens2_bf16_linear_uses_fused_bias_epilogue` passed in parallel only when it happened to run on the CPU. Once the MHR test passes, it holds its guard longer, which also exposed `moge3_refiner_forced_cache_pressure_preserves_full_unet_output` (12 of 24 parallel runs) and the batched-decode wrapper tests. Skipping the MHR test alone took that to 0 of 24. ## Changes - New `src/vision/addmm_oracle_test_support.rs` (test-only) documents the routing and provides `ROUTING_ROWS`, `routing_rows` and `single_row_addmm_rounds_twice` (Metal GPU default device). - The DINOv2 Objects linear, attention (qkv and proj), MLP (fc1 and fc2), SAM 3, SAM 3D Body and Sapiens2 linear oracles now run on two identical rows. There, the fused value is asserted on every backend and the separate value still differs, so the tests still prove the layer routes through addmm. The single-row value is also pinned per backend: separate on Metal GPU, and fused (the original authenticated oracle) elsewhere. No assertion was loosened to a tolerance. - The Sapiens2 RoPE oracle expects the ROCm capture on ROCm GPU and upstream's eager-equal value elsewhere. - The MHR geometry oracle stays bit-exact on Linux and for every non-transcendental field. Only on non-Linux hosts are the four transcendental fields bounded to 2 ULP. - These tests now hold `lock_default_device()` (or `mlx_test_guard()`): every test touched above, the tanh GELU oracle and feed-forward tests, the MoGe3 focal solver oracle, the MoGe3 refiner cache-pressure test, and the three dense VLM batched-decode wrapper tests. The lock was not added inside `batched_decode_harness`, because `vlm_family_batched_decode_tests` already calls it under `mlx_test_guard()` and the lock is not reentrant. - `tests/fixtures/sapiens2/README.md` now says the compiled RoPE capture is ROCm-only and gives the Metal and CPU value. ## Changes during review - Implementation and security reviews found nothing at CRITICAL or HIGH. A follow-up commit updates comments only: the helper module now says the CUDA addmm key comes from reading its source and was not measured, and it warns that Metal also rounds twice for a one-column output. It also adds the README note above. ## Validation Done on an Apple M1 Ultra, release build with `--features metal,accelerate`: - `cargo test --release --features metal,accelerate --lib vision::dense`: 192 passed, 3 of 3 runs in default mode and 3 of 3 with `--test-threads=1`. A stress loop of the same binary passed 40 of 40 default-mode runs (main failed 5 per run). - `cargo test --release --features metal,accelerate --lib dinov2_layers::policy`: 5 passed, 3 of 3 runs in each mode (main failed 5 of 5). - `vision::` filter, one default-mode run: 882 passed. `vision::encoders::dinov2`: 20 of 20 runs passed. - `cargo fmt --check`, `cargo clippy --features metal,accelerate --lib --tests -- -D warnings`, `scripts/ci/check_cross_repo_refs.py`. No real-checkpoint run was needed: no production code changed. The diff touches only `#[cfg(test)]` code and test files, and the production layers keep calling `mlxcel_core::addmm` exactly as upstream `nn.Linear` does on each backend. Not verified: - ROCm gfx1151. The two-row fused expectation on ROCm comes from reading the overlay's `AddMM::eval_gpu` (one rocBLAS or naive GEMM with beta for every M), not from a run. - Metal M3 or newer, where two rows select `gemv_wide`. That this path rounds once comes from its source (`alpha * result + beta * bias` computed in F32, then one cast), not from a run. - Linux on architectures other than x86_64 for the MHR oracle, and CUDA for the RoPE oracle, which has no capture. Closes #2626
7 of 8 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Add SAM 3D Objects reconstruction to
mlxcel predict: RGB plus mask or RGBA input produces object-local OBJ meshes, Gaussian PLY files, pose, and sparse latents. The implementation includes both flow stages, sparse decoders, FlexiCubes, shared DINO conditioning, bundled MoGe depth, and trained point-map dropout.Preserve the released BF16 main/F32 depth policy and validate all 3,895 checkpoint tensor names, shapes, and dtypes. Objects-specific interpolation, arithmetic, and compiled focal recovery match upstream while preserving other families’ defaults. Checked allocation policies retain complete supported outputs. A reviewed ROCm LayerNorm barrier closes a shared-memory ordering race.
Validation on ROCm gfx1151 with authenticated upstream, checkpoint, source, runtime, and process identities:
No parity thresholds, sampling defaults, model resolutions, or checkpoint dtypes were relaxed. Full reference arrays and historical failed runs remain outside git; compact provenance and reproducibility tools accompany the implementation.
Closes #2381