Skip to content

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

Description

@inureyes

Parent and dependency

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 test(xla): validate a LLaVA reference architecture end to end #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.

Activity

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

Metadata

Metadata

Assignees

No one assigned

    Labels

    area:inferenceGeneration, sampling, decoding (incl. speculative, DRY)area:modelsModel architectures, weights, loading, metadatapriority:lowLow prioritystatus:reviewUnder reviewtype:enhancementNew features, capabilities, or significant additions

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions