feat: add Youtu-VL XLA vision and MLA inference (#872) - #919
Conversation
206566c to
7886a65
Compare
|
Rebased onto Post-rebase validation:
The runtime now uses the explicit |
7886a65 to
e36cb34
Compare
Add checked static patch buckets, window/full attention isolation, 2D vision RoPE, the built-in merger, and resident IREE checkpoint loading for Youtu-VL. Map the dense DeepSeek-style MLA weights and compressed KV cache into token prefill, embedding prefill, single decode, and ragged continuous batching while retaining traditional interleaved RoPE. Wire filtered host preprocessing, placeholder selection, sequence export, and strict no-MLX-fallback capability routing; fix window-order restoration and preserve pre-allocation patch caps. Validated with xla-iree check, focused unit tests, clippy, pinned MLA IREE graph compilation, and one pinned checkpoint IREE vision execution. Closes #872
Fail closed when a Youtu-VL MLA checkpoint requests attention projection biases because the emitter does not yet load or apply the q_a, kv_a, and output bias tensors. Strengthen the vision restoration regression with a non-self-inverse window permutation so the prior incorrect permutation cannot satisfy the test accidentally. Validation: cargo fmt --all -- --check; cargo test -p mlxcel-xla --lib youtu_vl_selects_dense_mla_and_interleaved_rope; cargo test --lib reverse_indices_restore_original_group_order. Refs #872
Expose diagnostics-only eager and IREE checkpoints for patch projection, window/full layers, merger window order, and restored output. Add a deterministic two-layer dense MLA fixture for position zero, nonzero positions, replaced image embeddings, and padded IREE cache mapping. Validation: cargo test --lib youtu_vl; cargo test -p mlxcel-xla --lib youtu_vl; cargo clippy --lib --tests --no-deps; cargo check -p mlxcel-xla --features diagnostics. Refs #872
Keep the rebased Youtu-VL runtime on the explicit legacy artifact path until the full numeric contract is qualified.\n\nRefs #872
e36cb34 to
9171492
Compare
Expand one logical image or video placeholder per spatial shape to the exact number of merged visual tokens expected by Youtu-VL. Preserve already-expanded prompts and reject ambiguous partial expansion before runtime upload. Refs #872
Resolve dense checkpoint tensors by scanning every shard named by the index instead of trusting an individual stale weight-map entry. Continue rejecting missing, duplicate, and mixed canonical/wrapped namespaces. Refs #872
Match the pinned HF processor's spatial-merge grouping and channel-fast patch layout, and avoid regrouping those rows again at the IREE boundary. Correct the independent oracle reconstruction and add fail-fast preprocessing and selected vision-stage comparisons with partial diagnostic artifacts. Refs #872
Materialize the pinned checkpoint's BF16 vision boundaries while reusing the shared numeric LayerNorm, RMSNorm, GELU, and softmax decompositions. Keep the vision MLP's tanh GELU distinct from the merger's exact GELU. Refs #872
|
Published the bounded Youtu-VL follow-up at
Validation:
The real pinned checkpoint now reaches the unchanged intermediate gate quickly. An explicit BF16 contraction-input experiment produced byte-identical IREE |
Summary
DeepSeek-style MLA, compressed KV layout, and interleaved RoPE.
window/full attention selection, output restoration, and the built-in merger.
patch layout and expand one logical placeholder per image to the exact merged
token count.
shard while rejecting missing, duplicate, and mixed namespaces.
numeric LayerNorm, RMSNorm, GELU, and softmax decompositions.
tencent/Youtu-VL-4B-Instruct@8d30a0e49662a1d628a472b12df264dbcd768753and a fail-fast production runner that writes partial diagnostic tensors
before stopping at the first unchanged-threshold failure.
production gates remain unresolved.
Validation
cargo test --features xla-iree youtu_vl_ --lib: 12 passed, 1 ignored(the ignored gate requires the production checkpoint/runtime environment).
python -m unittest test_youtu_vl_reference_oracle.py: 13 passed.cargo fmt --all --checkcargo clippy --fix --allow-dirty --features xla-diagnostics --example xla_youtu_vl_reference_checkgit diff --checkdiagnostics-only worker topology.
Actual oracle evidence
vision RoPE match the pinned HF capture.
atol=0.02, rtol=0.02gate(
max_absolute=0.015625).max_absolute=0.03125with relative tolerance).layer.7.full: flat index 310,actual
0.3125, reference0.28515625, absolute0.02734375,max_absolute=0.203125, and 2,999 unchanged-threshold mismatches.within
0.015625of the HF reference, while the IREE layer-0 output reaches0.03125. The row-wise LayerNorm decomposition matches HF exactly; the firstdrift is the backend BF16 projection/attention contraction schedule.
so that ad hoc workaround was removed. No tolerance was changed.
retains every selected vision tensor for analysis instead of spending the
prior full-run duration on already-invalid language output.
Draft limitation
The current local CPU IREE target cannot establish production-reference BF16
contraction parity for the pinned HF eager schedule. Per #932, this is recorded
as an unresolved backend numeric contract, not mathematical equivalence.
Language logits/KV, greedy tokens, reset/reuse, and mixed continuous batching
therefore remain unqualified. This PR stays draft and the capability stays
fail-closed until the unchanged intermediate and token-exact gates pass on a
reproducible production-relevant contract.
Closes #872