feat(xla): support Molmo2 indexed attention pooling - #916
Conversation
inureyes
added
area:inference
207279c to
68eb41c
Compare
24cc11a to
8b9f917
Compare
|
Rebased onto Post-rebase validation:
The Molmo2 runtime now uses the explicit |
Emit and load the pinned Molmo2 flattened-patch ViT, indexed attention pooler, and additive projector through IREE. Extend the shared XLA text loader for Molmo2 fused weights and wire the owned prepared-prefill producer into the existing CLI and server lifecycle.
Molmo2 declares 27 logical ViT layers but persists only the 25-layer prefix needed by adapter selections `[-3, -9]`. Resolving the negative indices after truncation selected `[22, 16]` instead of the canonical `[24, 18]`. The eager loader and XLA emitter now resolve configured indices against the declared depth first, derive execution depth from the largest selected layer, and preserve configured order across negative and positive indices. Validation: `cargo test -p mlxcel-xla molmo2 --lib`; `cargo test -p mlxcel molmo2 --lib`; `cargo check -p mlxcel-xla`; `cargo check -p mlxcel --lib`; `cargo fmt --all -- --check`; `git diff --check`. Refs #871
Align the eager Molmo2 ViT with the checkpoint's declared gelu_pytorch_tanh activation instead of the global exact-erf helper, and reject incompatible activation metadata in the StableHLO contract. Add diagnostics-only MLX and IREE stage outputs for patch and position embeddings, the first and configured ViT blocks, selected-layer concatenation, masked gather and valid counts, pooling query and output, and the complete SwiGLU projection. The ignored checkpoint gate now compares those stages in order and reports the first divergence while preserving negative fixtures for layer order, denominator, and clamped-index leakage. Validation: cargo fmt --all -- --check; cargo check --features xla-iree --lib; cargo check --no-default-features --features xla-diagnostics-cpu --test molmo2_xla_vision_parity; cargo test -p mlxcel-xla --features diagnostics molmo2 --lib; cargo test --no-default-features --features xla-diagnostics-cpu --test molmo2_xla_vision_parity synthetic; cargo test -p mlxcel molmo2_uses_the_checkpoint_pytorch_tanh_gelu --lib; focused cargo clippy with the repository's pre-existing warning classes allowed. Refs #871
Prefer authoritative top-level vit_config and adapter_config sections over the empty vision_config compatibility object shipped by the pinned Molmo2 checkpoint. This keeps the eager MLX oracle aligned with IREE by resolving [-3, -9] against the declared 27-layer depth as [24, 18] before deriving the 25-block execution prefix. Add a regression fixture for the exact empty-wrapper checkpoint layout and preserve post-layer selection and configured concatenation order. Validation: cargo fmt --all -- --check; cargo test -p mlxcel molmo2_ --lib; cargo check --no-default-features --features xla-diagnostics-cpu --test molmo2_xla_vision_parity; cargo test -p mlxcel-xla --features diagnostics molmo2 --lib; git diff --check. Refs #871
The post-selected-layer actual gate passes pooling output but first diverges at projector.output_all with max_abs 0.1640625. The checkpoint stores all three projector weights as F32, and eager MLX and StableHLO agree on shapes and w1/w3/product/w2 ordering; the remaining semantic drift is that StableHLO reassociated MLX's x * sigmoid(x) into x / (1 + exp(-x)), changing F32 rounding before the wide w2 projection. Compute the reciprocal sigmoid first and multiply by x, matching the eager MLX operation graph. Add a mutation-sensitive emitter regression that rejects the quotient form and pins the sigmoid-then-multiply data dependency. Validation: - cargo test -p mlxcel-xla --features diagnostics molmo2 --lib -- --nocapture (19 passed) - cargo check --no-default-features --features xla-diagnostics-cpu --test molmo2_xla_vision_parity - focused mlxcel-xla Clippy with unrelated baseline lints allowed - cargo fmt --all -- --check - git diff --check The heavyweight actual-checkpoint gate was not repeated after this fix. Refs #871
The ignored actual-checkpoint reference check only emitted comparison output after both eager MLX and IREE work completed, so long checkpoint loading, projection, compilation, or invocation could appear stalled without any flushed progress. Emit explicitly flushed start/completion messages and 60-second heartbeats around MLX vision-only loading and projection plus IREE diagnostic compile/load and invocation. Flush every stage comparison before enforcing the existing limits, and add a synthetic regression for output flushing and the sub-five-minute heartbeat interval. No numerical threshold, model operation, or CSV changes are included. The heavyweight actual-checkpoint gate was not run. Validation: - Molmo2 diagnostics integration tests: 5 passed, 1 ignored - xla-diagnostics-cpu integration check - focused Clippy with existing repository baseline lints allowed - cargo fmt --all -- --check - git diff --check Refs #871
The reciprocal-before-multiply SiLU change reached the actual local-task VMFB and survived IREE lowering, but the real-checkpoint parity gate still failed at the same final projector boundary. Capture the projector w1, SiLU, w3, and product intermediates in both eager MLX and IREE diagnostics so the next checkpoint run identifies the first divergent operation without another speculative fix. Refs #871
Compare layer 18 before layer 24, then capture only the known failing row across the seven layer-24 block boundaries. Keep the production graph, strict thresholds, and checkpoint execution unchanged.
Keep the rebased Molmo2 runtime on the explicit legacy artifact path until the selected-layer numeric divergence is resolved.\n\nRefs #871
Use the merged contract-sensitive numeric helpers for Molmo2 tanh GELU, SiLU, and stable softmax so production emission and bounded probes cannot drift. Preserve the existing StableHLO operation order, including projector x * sigmoid(x); keep the rank-3 LayerNorm local because the shared helper is intentionally row-wise. Validation: Molmo2 diagnostics unit tests (21 passed), shared numeric helper tests (13 passed), synthetic Molmo2 parity tests (5 passed), rustfmt check, and git diff check. Refs #871
8b9f917 to
684c6ee
Compare
Rebase and bounded local-task evidenceThe feature branch is rebased onto The one authorized actual-checkpoint gate ran on rebased head Timings were: build 3m41s, checkpoint load 5s, CPU MLX eager projection 625s, IREE diagnostic compile/load 116s, IREE invocation 60s, and test execution 810.93s. Every staged intermediate remained inside the strict The exact first failing boundary is the final F32 dense projector: This local-task result does not justify changing attention, LayerNorm, activation order, or thresholds, and it does not qualify the artifact. No second actual-checkpoint or full-model run was attempted. Follow-up commit Focused validation after that integration:
The PR remains draft. The independent intermediate and token-exact acceptance gates are not complete, so this branch must not be merged yet. |
The Molmo2 vision parity gate created its IREE diagnostic projector without first applying the diagnostics-only local-task thread configuration added in #945. `configure_diagnostic_local_task_threads` had no callers on this branch, so the gate aborted before any comparison ran: xla_aux_create failed (status 13): iree/base/threading/thread_pthreads.c:159: INTERNAL; thread creation failed with 22 IREE parses its process-global flag registry when the first instance is created, so the call has to precede the projector load to have any effect. This is the same defect fixed for the Molmo v1 gate. Refs #871
The ViT and scatter parity gate passes at
|
MLX_ENABLE_TF32 |
vit.block.0 max_abs |
result |
|---|---|---|
1 (MLX default) |
0.0028572083, rms 0.0001811348 | FAILED at the first block |
0 (declared policy) |
0.000009536743, rms 0.00000040039913 | passed the full gate |
A 300x difference at block 0. The previously recorded first strict CUDA failure at selected
layer 18, and the projector.output_all failure at max_abs=0.1640625, rms=0.0026185848,
are both artifacts of the TF32 reference rather than properties of the emitted graph.
Stage results under the declared policy
vit.patch_embedding max_abs=0 rms=0
vit.position_embedding max_abs=0 rms=0
vit.positioned_embedding max_abs=0 rms=0
vit.block.0 max_abs=0.000009536743 rms=0.00000040039913
vit.probe.18.row.513.input max_abs=0.00048828125 rms=0.00001495746
vit.probe.18.row.513.attention max_abs=0.000000834465 rms=0.00000017476373
vit.probe.18.row.513.output max_abs=0.00048828125 rms=0.0000149557045
vit.selected.18 max_abs=0.00048828125 rms=0.0000020188866
vit.selected.24 max_abs=0.00048828125 rms=0.0000040258137
vit.concatenated max_abs=0.00048828125 rms=0.0000031845784
pool.gathered_masked max_abs=0.00048828125 rms=0.0000030708436
pool.valid_counts max_abs=0 rms=0
pool.output max_abs=0.00008392334 rms=0.0000032640364
projector.w1 max_abs=0.00007247925 rms=0.0000036522888
projector.silu max_abs=0.00007247925 rms=0.0000018101092
projector.w3 max_abs=0.00006866455 rms=0.0000035283178
projector.product max_abs=0.005859375 rms=0.000033962566
projector.output_all max_abs=0.046875 rms=0.00069711613
scatter-added embeddings max_abs=0.046875 rms=0.00049262145
The three embedding stages and pool.valid_counts are bit-exact.
Scope
This clears the ordered ViT intermediates, pooling, projector and scatter merge. The separate
logits, all-layer KV, token-exactness, CLI/server and lifecycle gates have not been run, so
the PR stays draft and the capability stays fail-closed until those pass.
The reference-policy defect and two build.rs architecture defects are filed on #932, since
they affect every family gate rather than this one.
Production path gates at
|
| check | result |
|---|---|
| text-only request on the multimodal XLA server | finish_reason=stop, content Red (see note below) |
| malformed base64 image | HTTP 400, OpenXLA image resolution cardinality mismatch: 1 image input(s) declared, 0 raw payload(s) resolved; refusing text fallback |
| capacity overflow | rejected before device execution: request exceeds the OpenXLA context capacity: effective_prompt_len=439 + max_new_tokens=24 > context_capacity=256 |
| two concurrent image requests | both complete, byte-identical content at temperature=0 |
/slots after completion |
all 4 slots back to idle, is_processing=false |
/health |
status=ok, max_batch_size=4, observability populated |
The malformed-image path is the behaviour this epic wants: it refuses a silent text fallback
and names the cardinality mismatch. The capacity rejection likewise happens on the host before
any device work.
Two observations worth a maintainer decision
1. The EOS token appears in returned content on this path. The text-only request returned
Red<|im_end|> with finish_reason=stop. tokenizer_config.json gives
eos_token = <|im_end|>, so the stop token is being detokenized into the response rather than
consumed. The same prompt on main's eager CLI returns Red with no trailing token. I have
not isolated whether this is XLA-versus-eager or CLI-versus-server, so I am reporting it as an
observation rather than attributing it; it needs one more bounded comparison to pin down.
2. The default context capacity rejects every Molmo2 image request.
MLXCEL_XLA_CONTEXT_CAPACITY defaults to 256, while a single fixture image expands to 439
prompt tokens, so the default configuration fails before generating a token. All results above
used MLXCEL_XLA_CONTEXT_CAPACITY=1024. Since the capacity is a static graph shape and part
of the compiled-artifact identity, this is a deliberate choice rather than a bug, but a
Molmo2 user on defaults hits a hard error, and the message does not name the variable to set.
Remaining before ready
Still outstanding, and none of these exist as tests on this branch:
- prefill logits and selected all-layer KV against the issue contract
- deterministic greedy output token-exact against the MLX reference (the parity gate compares
vision intermediates and the scatter merge, not decoded tokens) - request cancellation mid-generation
- capability predicate advertising only the loaded bundle and runtime conditions
- this head does not yet contain latest
main
So the PR stays draft. The intermediate, CLI, server and lifecycle evidence above is what has
been established.
Every OpenXLA path handed the terminating EOS id to its consumers as if it were generated output, while the eager MLX paths drop it. On GB10 the eager CLI returns `White` (1 token) and the XLA CLI returned `White<|im_end|>` (2 tokens); the XLA server showed the same with finish_reason=stop and completion_tokens one too high. The eager BatchScheduler tests merged_eos before recording or detokenizing, so the id never reaches generated_tokens or the incremental detokenizer. Every XLA path recorded or emitted first and tested after: XlaBatchEngine::pump pushed EngineEvent::Token before computing finish_reason at both the admission and decode sites, and the four single-sequence greedy loops plus XlaReferenceEngine::generate pushed into their output vector before the EOS test. XlaServeWorker counts one generated token per Token event and detokenizes whatever it receives, so the same defect produced both the leaked text and the inflated count. Introduce commit_token, which returns both whether a sampled token is client-visible and why the sequence ended, and route both pump sites through it. Collapse the four greedy loops onto one drive_greedy_loop so the contract is stated once, and keep that loop free of the session so it is unit-testable without a device. A token that ends generation by reaching the cap is still real output: it is collected and it is still handed to the streaming callback, which is why the callback is invoked before the budget is tested rather than after. finish_reason keeps its existing precedence of EOS over Length. The decode-width metric now counts a Stop finish as the one sampled token whose id was withheld, so suppressing the id does not silently drop a decode step from /metrics. Closes #963 Refs #566, #932, #916
Every OpenXLA path handed the terminating EOS id to its consumers as if it were generated output, while the eager MLX paths drop it. On GB10 the eager CLI returns `White` (1 token) and the XLA CLI returned `White<|im_end|>` (2 tokens); the XLA server showed the same with finish_reason=stop and completion_tokens one too high. The eager BatchScheduler tests merged_eos before recording or detokenizing, so the id never reaches generated_tokens or the incremental detokenizer. Every XLA path recorded or emitted first and tested after: XlaBatchEngine::pump pushed EngineEvent::Token before computing finish_reason at both the admission and decode sites, and the four single-sequence greedy loops plus XlaReferenceEngine::generate pushed into their output vector before the EOS test. XlaServeWorker counts one generated token per Token event and detokenizes whatever it receives, so the same defect produced both the leaked text and the inflated count. Introduce commit_token, which returns both whether a sampled token is client-visible and why the sequence ended, and route both pump sites through it. Collapse the four greedy loops onto one drive_greedy_loop so the contract is stated once, and keep that loop free of the session so it is unit-testable without a device. A token that ends generation by reaching the cap is still real output: it is collected and it is still handed to the streaming callback, which is why the callback is invoked before the budget is tested rather than after. finish_reason keeps its existing precedence of EOS over Length. The decode-width metric now counts a Stop finish as the one sampled token whose id was withheld, so suppressing the id does not silently drop a decode step from /metrics. Closes #963 Refs #566, #932, #916
1 participant
Summary
[24, 18]for[-3, -9]vit_configandadapter_configsections over the emptyvision_configcompatibility object(max_abs=0.05, rms=0.01)contractValidation
cargo test -p mlxcel-xla --features diagnostics molmo2 --lib(21 passed)cargo check --no-default-features --features xla-diagnostics-cpu --test molmo2_xla_vision_paritycargo test --no-default-features --features xla-diagnostics-cpu --test molmo2_xla_vision_parity synthetic(5 passed)cargo test --no-default-features --features xla-diagnostics-cpu --lib diagnostic_probe_snapshots_the_row_containing_the_actual_failure(1 passed)cargo test -p mlxcel molmo2_uses_the_checkpoint_pytorch_tanh_gelu --lib(1 passed)cargo check -p mlxcel --libcargo fmt --all -- --checkgit diff --checkActual-checkpoint evidence
The pinned
/home/inureyes/models/molmo2-4bcheckpoint passes stale-index shard discovery and selective eager-reference loading. Earlier CPU/local-task evidence first failed at the projector, but that evidence did not establish the production CUDA boundary.The correctly configured CUDA actual gate at
6203a402completed in about 144 seconds. It passed patch embedding, position embedding, positioned embedding, and block 0. Because that diagnostic compared selected layer 24 before selected layer 18, its first observed strict failure was:Commit
efd63b11reordered the selected-layer check into encoder order and added layer-24 row probes. Its authorized CUDA run completed in 141.67 seconds. The early stages passed again, but the new leading selected-layer comparison failed before those layer-24 probes:The identical maximum coordinate proves the row-513 discrepancy already exists at selected layer 18. It invalidates layer 24 as the first-cause target; fail-fast correctly stopped before the layer-24 row probes. No result was hidden or treated as a layer-24 substage result.
The block-0 maximum is unrelated to that row:
Therefore the available block-0 maximum does not show the row-513 discrepancy accumulating from block 0. No threshold relaxation is proposed.
Precision audit
Static checkpoint and graph inspection found:
float32_attention=trueandattn_implementation="sdpa"fast::scaled_dot_product_attentionQK^T, multiply by the scale, apply softmax, then multiply by VThese are mathematically equivalent contracts but not identical CUDA reduction, exponential, and accumulation plans. The observed shape — low RMS with one large value after many blocks — remains compatible with accumulated backend numerical sensitivity, but the current evidence does not prove whether attention, LayerNorm, or the MLP first creates the row-513 discrepancy.
Replacing production eager SDPA with a materialized attention matrix would be a substantial memory/performance policy change and would conflict with the checkpoint's explicit SDPA request. Changing StableHLO to imitate an unspecified MLX CUDA kernel plan is likewise not justified without locating the first divergent producer boundary. No production change is included.
Corrected bounded row diagnostic
The failing index decomposes as:
Commit
24cc11admoves the seven diagnostics-only producer snapshots from layer 24 to the actual first failing selected layer 18. The comparison order is now:vit.selected.18vit.selected.24A regression test pins all seven
vit.probe.18.row.513.*stage names immediately beforevit.selected.18, so fail-fast cannot hide them again. The diagnostic artifact identity is bumped tofirst-divergence-v4-layer18-row513.The seven extra IREE outputs still transfer only
7 * 1,152F32 values. Eager probe rows remain producer-local, host-owned snapshots. The production graph, production output ABI, checkpoint contract, and strict thresholds are unchanged.No actual checkpoint was run at
24cc11ad.Remaining gates
24cc11adonly when another actual run is authorizedThe PR remains draft. No CSV was generated or changed.
Invalidated pre-fix evidence
Earlier
[22, 16]evidence is invalid because negative indices were resolved against a truncated depth and the eager loader treated an empty compatibility wrapper as authoritative. The canonical pair is[24, 18]. The reciprocal/multiply-only projector hypothesis is insufficient. The6203a402layer-24 observation was valid but not the first encoder-order failure; theefd63b11run places the first currently observed strict CUDA failure at selected layer 18.Closes #871