feat: add Qwen2.5-VL OpenXLA vision path

[inureyes]

Summary

• add the Qwen2.5-VL OpenXLA/IREE vision graph with checkpoint-compatible window partitioning, patch/RoPE permutation, full/window attention scheduling, SwiGLU blocks, inverse restoration, and the family-specific patch merger
• keep text quantization separate from the unquantized vision tower when vision_config.skip_vision=true, while preserving the existing quantized Qwen2-VL path
• resolve the CUDA precision contract to F16 contractions with F32 accumulation, normalization, softmax, and outputs
• replace the eager manual F16 RMSNorm path with MLX fast_rms_norm, preventing square overflow from collapsing non-zero rows to zero
• add a strict actual-checkpoint oracle with ordered vision substages, M-RoPE, final logits/top-1, mutation controls, persisted JSON reports, and explicit multi-media/padded-bucket mask audits
• localize the remaining multi-media drift at full-attention layer 23 with Q/K/V, pre-projection context, projection, residual, norm, and MLP diagnostic seams
• retain the layer 16/17 boundary outputs, prove that inherited norm2 drift—not projection or weight semantics—causes the amplified layer-17 outliers, and reject an unsafe coarse F16 state-rounding hypothesis without weakening the oracle

Validation

• cargo fmt --all -- --check
• git diff --check
• cargo test -p mlxcel-xla --features diagnostics qwen25 -- --nocapture (9 passed)
• cargo test -p mlxcel-xla --features diagnostics diagnostic_layers -- --nocapture (1 passed)
• focused eager Qwen2.5-VL GPU/RMSNorm tests (3 passed)
• cargo check --features xla-diagnostics --example xla_qwen25_vl_window_check
• cargo clippy -p mlxcel-xla --features diagnostics --tests --no-deps (passed with existing unrelated warnings)
• earlier Qwen2-VL emitter, M-RoPE, root library, and IREE feature checks remain green

Review hardening and actual-checkpoint evidence

The first pinned CUDA run exposed a startup blocker before native execution: the root 4-bit text quantization contract was inherited by an unquantized Qwen2.5-VL vision tower whose intermediate_size=3420 is not divisible by the text group_size=64. Commit 5a6ee4b7 separates the resident vision contract while preserving text-embedding quantization.

Source and actual-oracle comparison then showed that the Qwen vision emitter bypassed the resolved device precision and built every contraction in F32, while the eager checkpoint path uses the F16 text-embedding dtype. Commit 592f96bb applies the CUDA F16-contraction/F32-sensitive-math policy to production and diagnostic graphs and fingerprints that policy in the compiled artifact identity.

The subsequent single-224-image investigation initially appeared to diverge late in the network, but host-owned producer snapshots proved that eager VisionRMSNorm itself collapsed non-zero F16 rows to zero. The manual fallback squared values in F16; magnitudes above sqrt(65504) ~= 255.94 overflowed before the mean, making the reciprocal RMS zero. Commits 31b05e91 through 230e5be4 progressively ruled out graph donation and added a focused characterization test.

Commit 52313107 applies the production fix: eager Qwen2.5-VL now uses MLX fast_rms_norm, matching upstream nn.RMSNorm behavior and the IREE graph's F32-sensitive normalization contract. The three focused GPU/RMSNorm tests pass. With that fix, the bounded single-224-image actual run passed every positive comparison under the unchanged strict atol=0.08, rtol=0.08 contract. Its Qwen2-style-full-attention mutation remained numerically within tolerance for that fixture, so numeric sensitivity alone is no longer treated as a valid structural guard.

A separate bounded actual run used two 140x140 checker images, each with grid_thw=[1,10,10]: 200 active patches in a 256-row bucket, exercising both non-identity window restoration and multiple media. Full/window outputs through layers 0, 7, and 15 passed. The first failure was full-attention layer 23:

• 729 of 256,000 values outside tolerance
• max_abs=101.30908
• error RMS 0.4572
• no non-finite values

Layer 24 onward propagated that divergence, so later-layer failures are not independent evidence. The report was persisted at /tmp/mlxcel-qwen25-vl-multiwindow-multimedia-report.json.

Static and focused runtime audits rule out host mask construction for this fixture:

• packed media boundaries are exactly [0,100,200]
• bucket size is 256, leaving 56 padded rows
• full and window masks differ structurally
• full attention has zero cross-media finite entries
• full attention has zero active-to-padding or padding-to-active finite entries
• each padded query exposes only its finite self key

Commit c6dccafe makes that structural audit the deterministic negative control and moves the next diagnostic seam to the first observed failure, layer 23. It exposes post-RoPE Q/K, V, the masked attention context before output projection, projected attention, residual, norm, and MLP on both runtimes. Diagnostic outputs validate the full bucket, including padding, before truncating to the 200 active patches. Production arithmetic, tolerances, and output contracts are unchanged.

The first two-media local-sync invocation at c6dccafe did not reach numerical comparison because the new diagnostic ABI exposed a host descriptor defect: xla_aux reported output 16 as rank 3/2, type 553648160/553648160, and elements 327680/327680. Output 16 is the layer-23 post-RoPE query. StableHLO correctly returns [256,16,80], while the Rust runtime had declared every non-merger result as [256,1280]; the element count matched but the meaningful head axes did not.

Commit 2d13e470 defines the complete ordered 25-output diagnostic descriptor schema and preserves rank-3 [patch,head,head_dim] descriptors for query, key, and value. The existing public comparison API still receives flat owned vectors only after the validated device transfer. A regression pins every output semantic, ordinal, rank, and shape for the actual 32-layer configuration, including output 16 query, output 17 key, output 18 value, and the final merger. No emitted graph or numerical operation changed.

The subsequent bounded two-media local-sync run at 2d13e470 completed the numerical comparison and validated that descriptor fix. Layer 15 still passed with zero failures (max_abs=0.18795204, error RMS 0.00795757), but the captured input to layer 23 was already divergent: 350 failures, max_abs=132.40136719, error RMS 0.49355285, and no non-finite values. Within layer 23, both the masked attention context and projected attention passed with zero failures (error RMS 0.00515040 and 0.00488825, respectively). This evidence moves the unresolved boundary before layer 23 rather than attributing it to that layer's attention implementation.

Commit 43f7da74 adds the compact follow-up capture: the seven ordered post-layer outputs for layers 16 through 22 appear before the layer-23 substages in both eager and IREE diagnostics. The descriptor regression now pins all 32 current outputs: rank-2 state outputs through layer 23 input/norm1, rank-3 query/key/value, the remaining rank-2 substages, and the merger. The strict report records the pre-probe layer plan and comparisons. Production arithmetic, tolerance, and output contracts remain unchanged.

The matching IREE 3.12 CUDA run at 43f7da74 completed and reduced the boundary to two adjacent layers. Layer 16 passed with zero failures (max_abs=0.48780060, error RMS 0.01032155); layer 17 was the first catastrophic failure with 392 values outside tolerance, max_abs=104.72387695, error RMS 0.41014538, and no non-finite values. The report is persisted at /tmp/mlxcel-qwen25-vl-43f7da74-cuda-report.json. Layers 18 onward are propagation and do not justify additional large captures.

Commit 93cc1e2f moves the ordered substage probe from layer 23 to layer 17 and retains only the layer 16 and 17 boundary outputs. The former pre_probe fields and labels are now generic diagnostic_layer names because the compact interval includes both the passing and first failing layers. The actual-config regression pins all 27 current output labels, ordinals, ranks, and shapes, including rank-3 layer-17 query/key/value descriptors. No duplicate layer output, production arithmetic, tolerance, or output contract was added.

The matching IREE 3.12 CUDA run at 93cc1e2f passes the layer-17 input, first normalization, masked attention context, projected attention, post-attention residual, and second normalization with zero failures. Query, key, and value have only 13, 8, and 1 marginal failures. The first catastrophic boundary is the MLP output: 338 failures, max_abs=104.05749512, error RMS 0.40832933, and no non-finite values. The report is persisted at /tmp/mlxcel-qwen25-vl-93cc1e2f-cuda-report.json.

Commit 9173a394 replaces that single opaque MLP result with the exact implemented SwiGLU seams: gate projection, SiLU gate activation, up projection, gated product, and down projection. The existing final MLP tensor is renamed to down projection rather than duplicated. The actual-config regression now pins all 31 output labels, ordinals, ranks, and shapes, including four [256,3420] intermediate outputs and the [256,1280] down projection. Runtime active-row truncation derives each flat vector length from its validated descriptor, preserving the wider intermediate axis. Production arithmetic, tolerance, CSVs, and production output contracts remain unchanged.

The bounded IREE 3.12 CUDA run at 9173a394 completed every diagnostic comparison, persisted /tmp/mlxcel-qwen25-vl-9173a394-cuda-report.json, and exited 101 only when the unchanged strict oracle panic enforced the recorded failures. Layer-17 input, norm1, attention context, projected attention, post-attention residual, and norm2 passed. The MLP substages showed:

• gate projection: 8 failures, max_abs=2.014036
• SiLU gate activation: 3 failures, max_abs=0.758430
• up projection: 23 failures, max_abs=4.825338
• gated product: 23 failures, max_abs=105.50305
• down projection: 338 failures, max_abs=104.057495

The gated product amplifies sparse projection outliers; it is not evidence that multiply itself introduced the first error. Source audit also corrected the quantization hypothesis for this fixture. Layer-17 gate/up/down weights and biases are dense F16 tensors with no affine scales or biases sidecars. Eager therefore selects UnifiedLinear::Regular, transposes [out,in], performs MLX's native F16 x @ W^T, then adds an F16 bias. The IREE loader widens the exact same F16 bytes to F32, the CUDA precision contract demotes only the activation and weight inputs back to F16 for dot_general, returns F32, and adds the widened F32 bias. Axes and operation order agree; the remaining concrete seam is native MLX F16 projection/result behavior versus IREE's F16-input/F32-result reduction and bias rounding. The current report does not prove which production side should change.

Commit 790c5dad adds only eager diagnostics for that decision. Using the same eager norm2 values and checkpoint weights, it computes dense F32 gate and up projections and compares each existing IREE output against both the native eager projection and the F32 control. No IREE output or descriptor is added. Every comparison now records the maximum-error flat index plus actual and expected values. One matching actual run can therefore determine whether IREE follows the native F16 result or the F32 accumulation control before any production correction is considered.

The bounded IREE 3.12 CUDA run at 790c5dad completed every comparison, persisted /tmp/mlxcel-qwen25-vl-790c5dad-cuda-report.json, and exited 101 only when the unchanged strict oracle panic enforced the recorded failures. Both controls retained exactly the same sparse failures and maximum-error indices:

• gate at flat index 159047: IREE -16.59533882, native eager F16 -18.60937500, eager dense F32 -18.61830139; both comparisons had 8 failures, with error RMS 0.00858201 and 0.00857863
• up at flat index 436067: IREE -12.20591164, native eager F16 -17.03125000, eager dense F32 -17.02601814; both comparisons had 23 failures, with error RMS 0.01666641 and 0.01664133

Native F16 and dense F32 eager results are close to each other while both remain far from IREE at the same outliers. This disproves projection result precision as the root cause for this fixture; changing production projection precision is not justified by the evidence.

Commit 89bed569 adds the next bounded diagnostic without changing the IREE ABI. It takes the already-captured IREE layer-17 norm2 values, rounds them through the same F16 input type used immediately before the CUDA dot_general, and replays the same checkpoint gate/up weights through the eager dense-F32 control. The two new host-side comparisons distinguish amplification of inherited norm2 drift from backend contraction/reduction behavior using the existing IREE gate/up outputs.

The bounded IREE 3.12 CUDA run at 89bed569 made that distinction definitive. Both projections replayed from captured IREE norm2 passed with zero failures:

• gate: max_abs=0.0001220703, error RMS 0.0000011104
• up: max_abs=0.0002899170, error RMS 0.0000013243

The same checkpoint weights and dense-F32 projection schedule reproduce IREE when supplied IREE's own norm2. Projection reduction, weight loading, axes, and bias semantics are therefore excluded; the sparse gate/up failures are inherited norm2 drift amplified by the wide projections and SwiGLU product.

A bounded source and runtime dtype audit then found an explicit contract mismatch. The actual checkpoint's patch weight and all twelve norm/QKV/projection/MLP weight-or-bias tensors in each block 0 through 17 are F16. Eager MLX preserves F16 at every block input, fast_rms_norm output, post-attention residual, second norm, and block output. In contrast, IREE Precision::F16 demoted only each dot_general input and returned F32, so patch state, normalization outputs, and both residual streams accumulated as F32 across layers.

Commit 2b1c1f62 tested a bounded coarse state-boundary hypothesis. It kept RMSNorm, softmax, RoPE, contraction accumulation, and the public ABI in F32 while rounding patch state, normalization results, post-attention residuals, and block outputs through F16.

The matching actual CUDA run rejected that hypothesis. Patch embedding still passed with max_abs=0.00024414, and layer 15 still passed with max_abs=0.125, but the first critical boundaries regressed relative to the 89bed569 baseline:

• layer 17: 392 failures at max_abs=104.7239 became 1,380 failures at max_abs=320
• down projection: 338 failures at max_abs=104.0575 became 1,504 failures at max_abs=319.2441
• merger: 170 failures at max_abs=0.8468 became 974 failures at max_abs=1.5832

The report is persisted at /tmp/mlxcel-qwen25-vl-2b1c1f62-cuda-report.json. Commit 53e0ffb1 reverts only the coarse patch, normalization, and residual state round trips, restoring those production files exactly to 89bed569 while preserving all diagnostics. The evidence requires an operation-by-operation backend rounding contract before another production correction is safe; no tolerance waiver is justified.

Remaining production validation

The earlier non-diagnostic pinned CUDA/IREE smoke completed successfully (Orange, one token, exit 0). Strict single-image positive acceptance passes. The multi-media diagnostic has localized inherited norm2 drift, but the rejected 2b1c1f62 experiment proves that applying blanket F16 round trips at visible state boundaries is not a safe correction.

The strict blocker is now a per-operator backend rounding contract that identifies where eager MLX rounds normalization, residual, bias, and intermediate results—not merely their observable tensor dtypes. That contract must be modeled and tested before another production change. Repeated multi-image execution, final logits/top-1 acceptance for the multi-media fixture, and mixed continuous-batch cancellation/slot-reuse validation also remain outstanding. This PR therefore remains draft.

No actual checkpoint was rerun after the revert commit. No tolerance, CSV, IREE output contract, or other production change is included.

Closes #866

[inureyes]

Rebased onto main@7fe1412d (including #923 and #934) and force-pushed as b45080d2.

Post-rebase structural validation:

• cargo fmt --all -- --check
• git diff --check
• cargo test -p mlxcel-xla --features iree --lib qwen25: 7 passed
• cargo check --features xla-diagnostics --example xla_qwen25_vl_window_check: passed
• cargo test -p mlxcel-xla --features diagnostics qwen25: blocked before the target tests by the existing gemma3n_qmv test-only Builder cfg/import defect already documented in this PR; the non-test diagnostic example builds

A fresh post-#923 two-image CUDA oracle was attempted with the existing local 4-bit checkpoint and fixture. Both the newly built binary and the preserved pre-rebase binary fail before numerical comparison while loading model.safetensors with an empty-device-vector vector::_M_range_check. Host evidence confirms this is an environment gate: nvidia-smi cannot communicate with the driver and /dev/nvidia* is absent, while the 3,073,720,461-byte checkpoint is intact (sha256:636982419c940321ac0f7793dc9dc3575a4ee4843a6b167b2e8c3d3cd25dacf4).

Therefore no post-#923 Q4 numeric result is claimed. The diagnostics artifact stays explicitly new_legacy_unqualified, the PR remains draft, and #866 remains blocked by #932 until CUDA access and the operator-level numeric contract are restored.

[inureyes]

Rebased this draft onto main@ec6fbb6d and force-with-lease updated the branch at
02d41a18.

The only conflict was the Qwen emitter import boundary. The resolution preserves
both the merged shared numeric_ops helpers and this branch's device
Precision selection. No family-local replacement of the shared numeric
foundation was retained.

Validation after rebase:

• cargo test -p mlxcel-xla --features diagnostics qwen25 --lib: 9 passed
• cargo fmt --all --check
• git diff --check

The PR remains draft. The production-reference CUDA Q4 probe required by #932
cannot be reproduced on this host, so no intermediate or token-exact claim is
being upgraded and no tolerance is changed.

[inureyes]

Current-head CUDA actual at 02d41a18, and a rejected state-rounding hypothesis

Run on the GB10 qualification host (driver 580.159.03, CUDA 13.0, IREE 3.12.0rc20260721,
Rust 1.93.1, MLX pin b7c3dd6d27f4). Working tree clean at 02d41a18.

Build prerequisite that invalidated the first attempts

The example could not execute MLX CUDA kernels at all until the build architecture was
pinned. src/lib/mlxcel-core/build.rs resolves the MLX architecture as
MLX_CUDA_ARCHITECTURES else detect_cuda_arch() else the literal "90a", and
detect_cuda_arch() shells out to nvidia-smi. In an environment where the GPU is not
exposed, the build silently falls back to 90a, which cannot run on this host's sm_121.

The symptoms are misleading and look like runtime or numeric faults:

• cudaGraphAddKernelNode(...) failed: invalid resource handle with CUDA graphs enabled
• cudaLaunchKernelExC(...) failed: no kernel image is available for execution on the device
  with MLX_USE_CUDA_GRAPHS=0

The live MLX build in this worktree (written 09:09) was 90a. After exporting
MLX_CUDA_ARCHITECTURES=121a and rebuilding, the runner completed normally. Recording this
because it is a build-environment failure, not a numeric one, and it should be separated
from numeric results.

Result: prior evidence reproduced exactly

Two 140x140 images, --context-capacity 256, unchanged atol=0.08, rtol=0.08. Report at
/var/tmp/mlxcel-cuda-qual/pr920/report.json. The comparison reproduces the earlier
89bed569 report value for value:

stage	max_abs	failures
vision.post_full_layer_23_capture_2	105.27258	586
vision.post_layer_17_diagnostic	104.72388	392
vision.merger_window_ordered	0.8468342	170

So the earlier numbers stand as current-head evidence. A separate single-image run at this
head passes every numeric comparison and then panics at the later negative control
(negative oracle accepted an identity patch permutation, line 883), consistent with the
already-documented weakness of the 224x224 fixture.

Hypothesis implemented and rejected: op-by-op f16 state rounding

Every MLX expected value in the report is exactly representable in float16 at every stage
(patch embedding, each layer output, norm1/norm2, query/key/value, attention context,
residual, gate, silu, up, gated product, down). Every IREE value is full f32. That is the
predicted signature of the contract gap: Precision::F16 demotes only dot_general inputs
and returns f32, while the eager path evaluates an f16 checkpoint in which every op
materializes its result at f16.

That was implemented completely, not coarsely: a Builder::round_state helper applied at
the patch embedding, norm1, norm2, QKV, post-RoPE query and key, the SDPA context, the
attention output projection, both residuals, gate, silu, up, the gated product and the down
projection. Softmax interior stayed f32 to match a fused SDPA, and the merger was excluded
because its outputs are not on the f16 grid.

Measured effect on the two-image fixture, most stages improved substantially:

stage	before	after
reordered_patch_embedding	0.003506	0.00024
post_full_layer_15	0.196769	0.125
post_layer_16_diagnostic	0.487801	0.35938
layer_17.norm2	0.251740	0.125
layer_17.query failures	13	1
layer_17.value failures	1	0

But the massive-activation channel regressed (post_full_layer_23 105.27 to 170.00,
failures 586 to 946), and decisively, the previously passing single-image case regressed
from all comparisons passing to 11 failing stages. The change is therefore rejected and not
included here. The patch is retained outside the repository as negative evidence at
/var/tmp/mlxcel-cuda-qual/pr920/rejected-f16-state-contract.patch.

This is the second refutation of the "round the state" family, after the earlier coarse
four-boundary attempt. Together they rule out intermediate storage dtype as the explanation,
which narrows the remaining search.

Where the divergence actually concentrates

Failures are not at the large-magnitude elements. At |expected| = 2088 the tolerance
atol + rtol*|expected| is about 167, wider than float16's own resolution there. The
failing elements are small-|expected| ones, and the ~1e33 max-relative figures come from
near-zero elements that are not themselves the failures.

The outliers localize to MLP intermediate channel 1727 and hidden channel 849, at the same
patch index within each of the two 100-patch images. The existing controls already showed
that replaying IREE's own norm2 through the dense f32 projection reproduces IREE's gate and
up outputs, so the projections are faithful and the difference enters at the normalization.
A sub-ULP difference in the dominant channel rescales every output of that row through the
RMS, which the wide projection and the SwiGLU product then amplify.

The blocker recorded in the PR body is unchanged, and the multi-media logits/top-1, repeated
multi-image, and mixed continuous-batch cancellation and slot-reuse gates remain outstanding.
This PR stays draft.

[inureyes]

TF32 reference policy checked here: immaterial, so the numbers above stand

I filed on #932 that MLX defaults MLX_ENABLE_TF32=1 and that
src/lib/mlxcel-xla/README.md declares that mode is not valid F32 reference evidence. That
defect decided the Molmo2 gate in #916 (fail to pass, 300x at vit.block.0) and moved the
Molmo v1 numbers in #918 by 12 to 27 percent, so it was necessary to re-check this PR before
relying on anything posted above.

Re-run at 02d41a18 with MLX_ENABLE_TF32=0, everything else identical
(/var/tmp/mlxcel-cuda-qual/pr920/report-notf32.json). The result is unchanged:

stage	TF32 default	MLX_ENABLE_TF32=0
vision.post_full_layer_23_capture_2	105.272583, 586 fail	105.272583, 586 fail
vision.post_layer_17_diagnostic	104.723877, 392 fail	104.723877, 392 fail
vision.layer_17.mlp_gated_product	105.503052, 23 fail	105.503052, 23 fail
vision.post_full_layer_31_capture_3	424.574951, 9273 fail	424.574951, 9273 fail
vision.merger_window_ordered	0.846834, 170 fail	0.846769, 171 fail

Every emitted-graph stage is identical to six decimal places. The only values that moved are
the host-side dense-f32 controls, which are genuine f32 host computations:

control	TF32 default	MLX_ENABLE_TF32=0
mlp_gate_projection_iree_norm2_dense_f32_control	0.000122	0.000076
mlp_up_projection_iree_norm2_dense_f32_control	0.000290	0.000122

That contrast is a clean confirmation of the mechanism: the flag reaches the f32 host controls
and does not reach this vision tower's contraction path, which runs at the checkpoint's
narrow dtype. So unlike #916 and #918, the divergence recorded here is a property of the
emitted graph and not of the reference policy.

The blocker stated in the PR body is therefore unchanged, and so is the conclusion from the
rejected op-by-op f16 state-rounding experiment described in my earlier comment.