Repository navigation
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 |
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 |
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
a4b641e to
7085564
Compare
7085564 to
8463378
Compare
The Qwen3-VL DeepStack test added in #921 indexes `packed_attention_bias[48 * 64 + 0]`, which `clippy::identity_op` rejects under `-D warnings`. Nightly verify run 38813575943 on main a426055 failed its `make verify-clippy` step on this line (`src/lib/mlxcel-xla/src/emitter/qwen3_vl.rs:1298`), while its test and video steps passed. The index is now `48 * 64`, the same element (row 48, column 0). Validation: `cargo clippy -p mlxcel-xla --all-targets -- -D warnings` and `cargo fmt --all -- --check` are clean, and the 7 `emitter::qwen3_vl` tests pass. Refs #921
Summary
Implements the OpenXLA/IREE vision encoder, DeepStack intermediate layer feature extraction/injection, and multi-axis M-RoPE language prefill path for
qwen3_vl(#868).Key Changes
src/lib/mlxcel-xla/src/vision.rs,src/lib/mlxcel-xla/src/qwen3_vl_runtime.rs):Qwen3VLVisionConfigwith bilinear position-embedding interpolation, 2D rotary embeddings across spatial blocks, main patch merger (use_postshuffle_norm = false), anddeepstack_visual_indexesintermediate mergers (use_postshuffle_norm = true).execute_vision_Sync) and returns both main visual features[N_vis, D]and per-branch DeepStack visual residuals[N_vis, D].src/lib/mlxcel-xla/src/lowering_prefill.rs,src/lib/mlxcel-xla/src/iree.rs,src/lib/mlxcel-xla/src/batch.rs):emit_prefill_with_deepstackandPreparedDeepStackPrefillwith dense[B, S, D]per-branch visual residual scatter aligned to<|image_pad|>token positions.XlaBucketSet::submit_deepstack_preparedand wiresXlaServingEngineandQwen3VlIreeHostPreprocessorfor both CLI and HTTP server multimodal requests.src/lib/mlxcel-xla/src/diagnostics/qwen3_vl_deepstack.rs,examples/xla_qwen3_vl_deepstack_check.rs):vision.patch_embedding,vision.position_embedding,vision.positioned_embedding,vision.block_0..23,vision.main_merger,vision.deepstack_merger_0..2, language prefillembedding,layer_0..27,final_norm,logits, same-input control, and negative controls (dropped branch, zeroed branch, shifted visual span).Qualification & Testing (
gfx1151ROCm,Qwen3-VL-2B-Instruct-4bit)xla_qwen3_vl_deepstack_check --device hip):tests/fixtures/test_image.pngandqwen25_test_1.png:max_abs <= 0.00416),vision.main_merger(max_abs = 0.000186), andvision.deepstack_merger_0..2(max_abs <= 0.0000839).same_input_control_failures: [](max_abs = 0.0).negative_dropped_branch_detected: true,negative_zeroed_branch_detected: true).greedy_top1_exact: true(final_top1 == eager_top1 == 1748ontest_image.png,16onqwen25_test_1.png),passed: true.tests/fixtures/test_image.png):"First state the color of this image in one word, then count from 1 to 10 separated by commas."(--max-tokens 16 --temp 0)MLXCEL_BACKEND=mlxCLI:"Orange, 1, 2, 3, 4, 5"MLXCEL_BACKEND=xlaCLI (--device hip):"Orange, 1, 2, 3, 4, 5"MLXCEL_BACKEND=xlaHTTP Server (/v1/chat/completions,--device hip):"Orange, 1, 2, 3, 4, 5"cargo test -p mlxcel-xla --lib --features diagnostics -- qwen3_vl deepstackcargo test --lib --features rocm,xla-iree -- qwen3_vl(42 passed)cargo test --lib --features rocm,xla-iree -- host_preprocessor(30 passed)cargo clippy --bin mlxcel --example xla_qwen3_vl_deepstack_check --features xla-diagnostics-rocm -- -D warningsCloses #868