Skip to content

feat: add Youtu-VL XLA vision and MLA inference (#872) - #919

Draft
inureyes wants to merge 14 commits into
mainfrom
feature/issue-872-youtu-vl-xla
Draft

feat: add Youtu-VL XLA vision and MLA inference (#872)#919
inureyes wants to merge 14 commits into
mainfrom
feature/issue-872-youtu-vl-xla

Conversation

@inureyes

@inureyes inureyes commented Jul 24, 2026

Copy link
Copy Markdown
Member

Summary

  • Add an explicit Youtu-VL XLA text configuration for the checkpoint's dense
    DeepSeek-style MLA, compressed KV layout, and interleaved RoPE.
  • Add bounded IREE vision buckets with checked grid planning, 2D vision RoPE,
    window/full attention selection, output restoration, and the built-in merger.
  • Match the pinned HF processor's spatial-merge-grouped, channel-fast flattened
    patch layout and expand one logical placeholder per image to the exact merged
    token count.
  • Recover tensors from stale safetensors index hints by scanning every named
    shard while rejecting missing, duplicate, and mixed namespaces.
  • Preserve explicit BF16 vision carrier boundaries while reusing the shared
    numeric LayerNorm, RMSNorm, GELU, and softmax decompositions.
  • Add an immutable independent HF-eager oracle pinned to
    tencent/Youtu-VL-4B-Instruct@8d30a0e49662a1d628a472b12df264dbcd768753
    and a fail-fast production runner that writes partial diagnostic tensors
    before stopping at the first unchanged-threshold failure.
  • Keep the family capability fail-closed while intermediate and token-exact
    production gates remain unresolved.

Validation

  • All 18 pinned checkpoint artifact hashes and the fixture hash matched.
  • cargo test --features xla-iree youtu_vl_ --lib: 12 passed, 1 ignored
    (the ignored gate requires the production checkpoint/runtime environment).
  • Focused HF patch-layout and stale-index resolver regressions passed.
  • python -m unittest test_youtu_vl_reference_oracle.py: 13 passed.
  • cargo fmt --all --check
  • cargo clippy --fix --allow-dirty --features xla-diagnostics --example xla_youtu_vl_reference_check
  • git diff --check
  • The pinned HF capture and local-task IREE run completed with bounded
    diagnostics-only worker topology.

Actual oracle evidence

  • Resized pixels, flattened patches, window ordering, spatial metadata, and
    vision RoPE match the pinned HF capture.
  • Patch projection passes the unchanged atol=0.02, rtol=0.02 gate
    (max_absolute=0.015625).
  • Layer 0 passes that gate (max_absolute=0.03125 with relative tolerance).
  • The first failing selected stage is layer.7.full: flat index 310,
    actual 0.3125, reference 0.28515625, absolute 0.02734375,
    max_absolute=0.203125, and 2,999 unchanged-threshold mismatches.
  • Feeding the IREE patch projection into the HF layer-0 implementation stays
    within 0.015625 of the HF reference, while the IREE layer-0 output reaches
    0.03125. The row-wise LayerNorm decomposition matches HF exactly; the first
    drift is the backend BF16 projection/attention contraction schedule.
  • Explicit BF16 contraction input casts produced byte-identical IREE results,
    so that ad hoc workaround was removed. No tolerance was changed.
  • The fail-fast runner now stops this case in roughly 39 seconds after build and
    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

@inureyes inureyes added status:review Under review type:enhancement New features, capabilities, or significant additions priority:low Low priority area:models Model architectures, weights, loading, metadata area:inference Generation, sampling, decoding (incl. speculative, DRY) labels Jul 24, 2026
@inureyes
inureyes force-pushed the feature/issue-872-youtu-vl-xla branch from 206566c to 7886a65 Compare July 27, 2026 00:02
@inureyes

Copy link
Copy Markdown
Member Author

Rebased onto main@7fe1412d after #934 and force-pushed as 7886a657.

Post-rebase validation:

  • cargo fmt --all -- --check
  • git diff --check
  • cargo test --lib youtu_vl: 24 passed
  • cargo test -p mlxcel-xla --features iree --lib youtu_vl: 9 passed, 2 intentionally ignored
  • real 27-layer actual-schedule IREE CPU compile gate: 1 passed

The runtime now uses the explicit new_legacy_unqualified auxiliary-artifact path. That is intentional: #934 is only the versioned numeric-contract foundation, while the complete pinned ~10 GB HF-reference-versus-production numeric report is still pending. This PR therefore remains draft and must not be merged yet; #872 remains blocked by #932.

inureyes added 10 commits July 27, 2026 11:22
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
@inureyes
inureyes force-pushed the feature/issue-872-youtu-vl-xla branch from e36cb34 to 9171492 Compare July 27, 2026 02:24
inureyes added 4 commits July 27, 2026 12:42
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
@inureyes

Copy link
Copy Markdown
Member Author

Published the bounded Youtu-VL follow-up at 20c2bffd:

  • exact logical placeholder expansion with fail-closed partial-expansion checks;
  • stale safetensors index recovery by scanning every shard named by the index;
  • pinned HF spatial-merge/channel-fast patch layout and corrected oracle reconstruction;
  • fail-fast preprocessing/selected-vision comparisons with partial diagnostic artifacts;
  • explicit BF16 vision carrier boundaries using the shared numeric helpers.

Validation:

  • Rust Youtu suite: 12 passed, 1 environment-gated ignored
  • Python oracle contract: 13 passed
  • focused patch-layout, shard-resolution, and emitter tests passed
  • fmt, clippy, and git diff --check passed

The real pinned checkpoint now reaches the unchanged intermediate gate quickly.
Preprocessing and metadata are exact; patch projection passes
(max_absolute=0.015625), and layer 0 passes. The first selected failure is
layer.7.full at flat index 310: actual 0.3125, HF 0.28515625,
absolute=0.02734375, max_absolute=0.203125, 2,999 threshold mismatches.

An explicit BF16 contraction-input experiment produced byte-identical IREE
results and was removed. The LayerNorm decomposition matches HF exactly; the
remaining drift begins at the backend BF16 projection/attention contraction
schedule. No tolerance changed. The PR remains draft and fail-closed; language
KV/logits/tokens and lifecycle gates are not claimed.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:inference Generation, sampling, decoding (incl. speculative, DRY) area:models Model architectures, weights, loading, metadata priority:low Low priority status:review Under review type:enhancement New features, capabilities, or significant additions

Projects

None yet

Development

Successfully merging this pull request may close these issues.

feat(xla): support Youtu-VL windowed vision and MLA inference

1 participant