feat(xla): support Youtu-VL windowed vision and MLA inference
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
## Parent and dependency
- Parent epic: #566
- Blocked by #862
- Existing MLX implementation: `src/loading/vlm_youtu_vl.rs`, `src/vision/youtu_vl.rs`, `src/vision/encoders/youtu_vl.rs`, `src/vision/processors/youtu_vl.rs`, and `src/models/youtu_vl_lm.rs`
## Problem
Youtu-VL combines a SigLIP2-style visual encoder with flattened Linear patches, variable spatial shapes, windowed attention, 2D vision RoPE, and a built-in merger that already projects to language hidden size. Its text backbone is a dense DeepSeek-V3-style MLA model with traditional/interleaved RoPE, not a Qwen M-RoPE model. The processor enforces a per-image patch cap before allocation.
Aliasing this family to either generic SigLIP/LLaVA or Qwen-VL would lose window/grid semantics or use the wrong language position model.
## Goal
Support a pinned Youtu-VL checkpoint end to end on XLA, including safe flattened-patch preprocessing, IREE windowed vision/merger execution, DeepSeek-style MLA text prefill/decode, placeholder replacement, CLI, and serving.
## Required implementation
- Reuse/share the processor's smart resize, SigLIP normalization, flattened `[total_patches, channels * patch_size^2]` layout, per-image `(h_patches, w_patches)` metadata, divisibility rules, checked allocation arithmetic, and `max_patches_per_image` enforcement.
- Define static patch/media buckets with actual counts and spatial shapes. Reject over-capacity images before allocating/uploading the flattened tensor.
- Emit Youtu patch projection, position/2D RoPE helpers, window index and cumulative boundaries, full/window block selection, normalization/MLP variants, and built-in patch merger through StableHLO/IREE.
- Preserve per-image/window attention isolation and output restoration order. The merger output is already in language hidden size; do not add another projector.
- Replace configured `image_token_id` positions with features. Preserve the current fallback to `video_token_id` only when no image token is present, and advertise video support only if its request/preprocessor semantics have an oracle fixture.
- Reuse the existing XLA DeepSeek-V3/MLA primitives only after mapping Youtu's dense/bias/traditional-interleaved-RoPE configuration and sanitized weights exactly. Do not select Qwen M-RoPE because the input is visual.
- Include patch/window/grid/merger config, placeholder IDs, MLA head decomposition, RoPE flags, and static limits in artifact compatibility identity.
- Integrate filtered loading, capabilities, CLI/server batching, cancellation, and text fallback without a duplicate full MLX decoder or silent MLX vision fallback.
## Non-goals
- General Qwen-VL M-RoPE.
- Unbounded dynamic spatial shapes.
- Additional connector layers after the built-in Youtu merger.
## Validation
- Pin one Youtu-VL checkpoint/revision and compare resized pixels, flattened patch ordering, spatial shapes, window indices/boundaries, vision RoPE, selected full/window block outputs, merged visual tokens, placeholder selection, MLA prefill logits/KV, and greedy tokens against MLX and an independent oracle.
- Cover multiple aspect ratios, maximum patch boundary, one-patch overflow rejected before allocation, multiple images, cross-image/window isolation, image-token and documented video-token fallback, no-media text, cancellation, and slot reuse.
- Add negative fixtures for CHW/patch flattening order, extra projection, Qwen position routing, and window restoration errors.
- Run CLI and mixed continuous-batch server tests on a production-relevant IREE target, plus #862 and text-only DeepSeek/MLA gates.
## Acceptance criteria
- [ ] Checked processor allocation and per-image patch limits are preserved before XLA upload.
- [ ] Windowed vision and built-in merger match intermediate oracles; no extra projector is applied.
- [ ] Youtu's dense MLA/traditional-interleaved-RoPE text contract is selected explicitly.
- [ ] Placeholder replacement and any advertised video fallback are independently validated.
- [ ] Deterministic output is token-exact on the pinned fixture.
- [ ] No full decoder duplication or silent MLX vision fallback remains.
- [ ] Existing text-only DeepSeek/MLA regressions remain green.
Contributor guide
Research direction
Start with the existing MLX implementation in src/loading/vlm_youtu_vl.rs, src/vision/youtu_vl.rs, src/vision/encoders/youtu_vl.rs, src/vision/processors/youtu_vl.rs, and src/models/youtu_vl_lm.rs, then review the blocked dependency #862. Validate against the pinned checkpoint and MLX/oracle intermediate outputs before running the CLI and mixed continuous-batch server tests. Done means token-exact output with the listed isolation, placeholder, allocation, cancellation, and text-only regression checks passing.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- ai, backend, machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 25/100