feat(xla): add Qwen3-VL DeepStack vision path - #921
Conversation
inureyes
added
status:review
396950c to
e0c9317
Compare
|
Rebased onto Post-rebase validation:
The runtime now uses the explicit |
Add a resident IREE Qwen3-VL vision emitter and runtime with learned-position interpolation, packed nonwindowed attention, the main merger, and every ordered DeepStack branch from the immutable checkpoint. Derive the dense nested text contract, preserve interleaved M-RoPE, route CLI and server image requests through owned DeepStack prefill payloads, and fail closed when an ordinary prefill would omit side features. Validate emitter/runtime fixtures, host and server routing boundaries, the XLA-IREE build and clippy paths, and one pinned CUDA/IREE image generation that returned `Orange`. Refs #868
Read the checkpoint pixel bounds and PIL bicubic resampler instead of using the historical Qwen2 defaults, and include the effective policy in the compiled artifact identity so incompatible processor artifacts cannot be reused. Qualify the 512-patch bucket required by two canonical 256x256 images, retain per-image packed attention isolation, cap the corresponding DeepStack payload at 128 merged positions, and cover processor, overflow, ABI, and context-capacity boundaries. Refs #868
Compare the eager MLX and resident IREE main vision merger and every ordered DeepStack branch from one qualified processor payload. Reuse the production DeepStack prefill loop to expose post-injection hidden states and final logits without changing serving outputs, and reject dropped or zeroed actual-checkpoint branches. Refs #868
Use the checkpoint's gelu_pytorch_tanh contract in the eager vision blocks while preserving exact GELU in the main and DeepStack patch mergers. Add eager numerical and emitted-graph regression coverage for the activation split.
Block-level comparisons can identify LayerNorm as the first failing stage even when its input already contains tolerated quantized-matmul drift. Re-run the eager norm1 and norm2 operators on the exact IREE block input and post-attention residual, then report dedicated same-input comparisons. Add deterministic eager and StableHLO contract tests for the centered-variance F32 LayerNorm formula. Validated with rustfmt, focused eager and emitter tests, and the xla-diagnostics example check. Refs #868
The same-input LayerNorm probe proves block 2 normalization is within tolerance when both runtimes consume identical activations, so the remaining divergence is accumulated upstream rather than a LayerNorm precision defect. Capture block 0 input, norm1, attention, residual, norm2, MLP, and output on both eager MLX and IREE diagnostic paths. The oracle now reports these stages before block 2 without changing production precision, tolerances, or checkpoint execution. Validated with the focused Qwen3-VL emitter tests, diagnostics library checks, the diagnostics-enabled oracle example check, and Clippy for mlxcel-xla. Refs #868
Qwen3-VL vision parity first crosses tolerance at block 2 after block 0 and block 1 remain within tolerance, leaving fused quantized projections and IREE's host-dequantized dense F32 weights as the primary unresolved arithmetic boundary. Add diagnostics-only same-input block 0 and block 1 controls that run the production fused-QMM path beside a dense F32 path using the exact affine host widening shared with the IREE checkpoint loader. Record both outputs in the actual-checkpoint oracle without changing production execution or tolerance policy. Validated with cargo fmt, the focused dense-F32 control unit test, cargo check for the xla-diagnostics example on the GB10/IREE CUDA setup, and cargo clippy for the same example. Refs #868
Keep the rebased Qwen3-VL runtime on the explicit legacy artifact path until its Q4 numeric oracle is rerun against the current CUDA kernel.\n\nRefs #868
e0c9317 to
a4b641e
Compare
|
Rebased this draft onto The only manual resolution was a formatting-only conflict in the DeepStack C Validation after rebase:
The PR remains draft. The post-#923 production-reference CUDA Q4 probe required |
Post-#923 current-head CUDA actual at
|
| stage | max_abs | failures |
|---|---|---|
vision.block_2 |
0.22236824 | 2 |
vision.block_9 |
0.31624985 | 31 |
vision.block_10 |
8.69421 | 169 |
vision.block_23 |
1078.1646 | 89017 |
The mergers all fail downstream of that: main_merger 6953, deepstack_merger_2.fc1 25624,
deepstack_merger_1 78. deepstack_merger_0 output still passes with 0 failures.
The substage evidence points at accumulation, not at a per-operator defect
Within block 2, the individual operators agree, while the normalizations fail only on
accumulated input:
| stage | max_abs | failures |
|---|---|---|
vision.block_2.input |
0.221995 | 0 |
vision.block_2.attention |
0.0225568 | 0 |
vision.block_2.mlp |
0.0696449 | 0 |
vision.block_2.norm1 |
0.27341986 | 1 |
vision.block_2.norm2 |
0.4035015 | 283 |
and the same-input controls pass outright:
| stage | max_abs | failures |
|---|---|---|
vision.block_2.norm1.same_input |
0.0849018 | 0 |
vision.block_2.norm2.same_input |
0.0423183 | 0 |
vision.block_0.output.same_input.fused_q |
0.111899 | 0 |
vision.block_1.output.same_input.host_dequant |
0.0266669 | 0 |
Given identical inputs, the graph reproduces the reference within tolerance at every probed
operator, including the fused Q4 path. The failures appear only once each block consumes the
previous block's drifted output, and the normalization is where a small input difference is
rescaled across the row and becomes visible.
This is the same amplifier identified on Qwen2.5-VL in #920, where replaying the reference's
own normalization output through the projection reproduced the candidate exactly. In both
families the failing elements are small-|expected| values with absolute differences near
0.2; the very large max-relative figures come from near-zero elements that are not the
failures themselves.
DeepStack language-layer injection, final-position logits, greedy token exactness, mixed
batching, cancellation and slot reuse were not reached, since the first strict intermediate
failure stands. This PR stays draft and the capability stays fail-closed.
TF32 reference policy checked here: immaterial, so the numbers above standI filed on #932 that MLX defaults Re-run at
Consistent with the mechanism: the flag reaches genuine f32 contractions only, and this So the divergence recorded here is a property of the emitted graph, not of the reference |
1 participant
Summary
Add dense Qwen3-VL image execution to the OpenXLA backend with a resident IREE vision graph that implements learned-position interpolation, packed nonwindowed attention, the main merger, and every configured DeepStack branch. Route CLI and continuous-batch server requests through compact per-layer DeepStack payloads, while failing closed if an ordinary prefill would omit those side features.
Honor the checkpoint processor contract by reading its smart-resize pixel bounds and PIL bicubic resampler, binding those values into the compiled artifact identity, and qualifying a 512-patch bucket for two canonical 256x256 images and 128 merged DeepStack positions.
Add diagnostics that drive eager MLX and IREE from the same production processor patch/grid payload. The ordered oracle compares patch projection, interpolated position embeddings, their sum, every vision block, the main merger and DeepStack merger internals, final merger outputs, hidden state after each target language layer's
deepstack_process, and final-position logits. Dropped and zeroed actual-checkpoint DeepStack branches remain negative controls.Align the eager Qwen3-VL transformer-block MLP with the checkpoint's
gelu_pytorch_tanhcontract. The IREE block path already used tanh GELU, while the main and DeepStack patch mergers correctly retain exact GELU.Capture blocks 0 and 2 at seven ordered boundaries on both eager MLX and IREE: input, norm1, attention, post-attention residual, norm2, MLP, and output. Numeric failures are accumulated so every vision comparison runs, and
--reportis persisted before the oracle exits.Add a same-input LayerNorm probe that re-runs eager block 2 norm1 and norm2 on the exact IREE block input and post-attention residual. This separates normalization arithmetic from upstream quantized-matmul drift without changing production arithmetic, comparison thresholds, or CSV files.
Add diagnostics-only same-input projection controls for blocks 0 and 1. Each control runs the production fused-QMM path beside a dense F32 path that uses the exact affine host widening shared with the IREE checkpoint loader, isolating projection accumulation order without changing production execution.
Related issues
Closes #868
Parent epic: #566
Type of change
feat— new user-visible featurefix— bug fixperf— performance improvementrefactor— internal restructuring without behavior changechore— build, CI, dependencies, release infrastructuredocs— documentation onlytest— tests onlyTest plan
cargo fmt --all -- --checkcargo check --features xla-ireecargo check --example xla_qwen3_vl_deepstack_check --features xla-diagnosticscargo clippy -p mlxcel-xla --features iree --fix --allow-dirty --allow-stagedcargo clippy --features xla-iree --lib(existing warnings in unrelated XLA/audio code)cargo clippy --example xla_qwen3_vl_deepstack_check --features xla-diagnostics(existing unrelated warnings)cargo test -p mlxcel-xla qwen3_vl--lib --features xla-ireeprocessor/prefill/server testsqwen3-vl-2bcheckpoint (Orange, 1 token, exit 0; reported generation time 28.13s)49f9a4f9: patch projection, position embeddings, positioned embeddings, block 0, and block 1 pass; the first failure isvision.block_2with only 2 failing values andmax_abs=0.22236824b10c81d1persisted/tmp/mlxcel-qwen3-vl-b10c81d1-report.jsonafter comparing every vision stageb10c81d1: input passes; norm1 has 1 failure; attention passes; post-attention residual has 1 failure; norm2 has 283 failures; MLP passes; output has 2 failures7a26a928completed in 29.5 seconds and persisted/tmp/mlxcel-qwen3-vl-7a26a928-report.json7a26a928: norm1 passes with 0 failures andmax_abs=0.084901; norm2 passes with 0 failures andmax_abs=0.042318396950c9persisted/tmp/mlxcel-qwen3-vl-396950c9-cuda-report.json396950c9: fused QMM passes withmax_abs=0.1118994; host-dequantized dense F32 passes and is closer withmax_abs=0.0191097396950c9: fused QMM passes withmax_abs=0.0400119; host-dequantized dense F32 passes and is closer withmax_abs=0.0266669396950c9remains the first failure with 2 failing values andmax_abs=0.22236824; both same-input block 2 LayerNorm probes passcargo test -p mlxcel-xla emitter::qwen3_vl::tests::diagnostic_bucket_appends(2 passed)cargo test -p mlxcel-xla emitter::qwen3_vl::tests::layer_norm_emits_centered_variance_f32_contractcargo test --lib vision::encoders::qwen3_vl::tests::block_layer_norm_matches_centered_variance_f32_contractrustfmtfor the five block 0 probe filescargo check -p mlxcel-xla --features diagnosticscargo clippy -p mlxcel-xla --features diagnostics --fix --allow-dirty --allow-staged(existing warnings only; unrelated automatic fixes were excluded)cargo deny check(CI)cargo clippy --all-targets -- -D warningscargo test --releaseNotes for reviewers
The
396950c9report proves the block 2 LayerNorm implementation is not the source of the tolerance failures. With ordinary accumulated inputs, norm1 has 1 failure atmax_abs=0.27341986and norm2 has 283 failures atmax_abs=0.40350151; with the exact IREE inputs fed to eager MLX, both operators pass with zero failures. Attention and MLP also pass in the ordinary run, while the block output returns to only 2 failures. The evidence therefore points to upstream accumulated drift.Static precision audit does not support an F16/BF16/F32 production fix. The Qwen3-VL emitter uses the default F32 builder and F32 inputs, weights, and contractions. Eager patch projection promotes F32 input and BF16 weights to F32, while the quantized blocks use MLX CUDA's F32
qmm_naivepath for this input shape. The remaining arithmetic seam is operator layout: eager MLX dequantizes Q4 weights on the fly inside fused QMM, whereas IREE consumes host-dequantized dense F32 weights in an F32 dot. Their reduction and dequantization order can differ even though their exposed dtype is the same.The exact-IREE-input block controls confirm that seam. For both block 0 and block 1, the IREE result is consistently closer to eager host-dequantized dense F32 than to eager fused QMM. This is a non-MoE dense architecture, but the local checkpoint configuration uses Q4 affine projections. The evidence identifies MLX fused Q4 QMM versus IREE host-dequantized dense F32 accumulation-order drift; it does not establish a production semantic bug.
Strict intermediate MLX tolerance remains blocked without either emulating MLX CUDA fused-QMM kernel ordering in the IREE path or explicitly changing the acceptance/reference policy. Neither change is made in this PR.
Position interpolation is the only confirmed dtype mismatch: eager MLX performs the interpolation in BF16 before adding it to the F32 patch projection, while IREE widens the table and interpolation weights to F32. That stage passes the actual oracle, so changing production interpolation precision would be speculative.
The pre-fix
Orangeresult was invalid canonical evidence because it used the wrong resize grid and resampler. After the processor-contract fix, a bounded CUDA/IREE run with the pinned checkpoint and fixture completed successfully, producedOrangeas one token, and exited 0.The unrestricted targeted
cargo test --features xla-iree ...command attempted to link unrelated integration binarypipeline_cli_real_modelsand failed in the platform IREE static runtime with__stack_chk_guard@@GLIBC_2.17; DSO missing from command line. The equivalent library-only focused tests pass. A diagnostics-feature test build also encounters an unrelated existinggemma3n_qmvtest-onlyBuilderimport error outside the Qwen3-VL path; diagnostics library and oracle example checks pass.Checklist
docs/if user-facing behavior or supported models changed// Used by: ...comments on any shared function modified.env, or CSV files committed