Skip to content

feat(xla): support Qwen2.5-VL windowed vision attention #866

Description

@inureyes

Parent and dependency

Problem

Qwen2.5-VL is not a checkpoint-name alias for Qwen2-VL. Its vision tower uses RMSNorm/SwiGLU variants and alternates configured full-attention blocks with windowed blocks. The implementation computes a window_index, window-specific cumulative sequence lengths, and a restoration order around the patch merger. Reusing Qwen2-VL's packed full-attention graph would be numerically wrong even when tensor shapes load successfully.

Goal

Extend the qualified Qwen2-VL XLA path to Qwen2.5-VL with exact window partitioning/reordering, vision-block semantics, processor behavior, M-RoPE state, and end-to-end validation.

Required implementation

  • Share the Qwen family processor, patch embedding, vision RoPE, artifact metadata, and M-RoPE language path where the existing MLX implementations are truly identical. Keep family-specific code/config explicit where semantics differ.
  • Emit Qwen2.5-VL VisionRMSNorm, QKV attention, SwiGLU MLP, configured full-attention blocks, and window-attention blocks through StableHLO/IREE.
  • Reproduce window_index, cu_window_seqlens, padding, spatial-merge-unit grouping, and inverse/restoration ordering from qwen2_5_vl.rs. Treat these as processor/graph inputs with validated static maxima or emit deterministic graph constants when fully shape-derived.
  • Ensure no token from one window/media segment attends outside its permitted range. Full-attention block indices must come from checkpoint configuration and be included in artifact compatibility identity.
  • Validate patch/temporal/merge/window sizes, block-index range and uniqueness, grid divisibility, reordered token count, and merger output count before native execution.
  • Integrate Qwen2.5-VL family detection and filtered weight loading without falling back to the Qwen2 vision graph on unknown/missing configuration.
  • Preserve the feat(xla): support Qwen2-VL as the representative M-RoPE VLM #865 CLI/server/cancellation/accounting path and feat(xla): support multimodal 3D RoPE positions and decode deltas #864 per-slot M-RoPE state.
  • Keep video support explicit: validate at least one temporal-grid fixture if advertised; otherwise reject it with a family-specific unsupported error.

Non-goals

  • Qwen3-VL fast position interpolation or deepstack feature injection.
  • General dynamic-shape window attention.
  • Reworking the already-qualified Qwen2-VL full-attention implementation.

Validation

  • Pin a Qwen2.5-VL checkpoint/revision and compare processor tensors, window indices, both cumulative-length forms, reordered patch embeddings, a windowed block, a full-attention block, restored/merged features, M-RoPE positions, logits/KV, and greedy tokens against the MLX path and an independent oracle.
  • Include fixtures where the grid spans multiple windows, contains padding at a window boundary, uses multiple media segments, and exercises every configured full-attention block.
  • Add a negative test that intentionally uses Qwen2 full attention and proves the intermediate oracle catches the divergence.
  • Run CLI and mixed continuous-batch server tests, including cancellation and slot reuse.
  • Execute on a production-relevant IREE target and run all feat(xla): support Qwen2-VL as the representative M-RoPE VLM #865/Qwen2/text-only regression gates.

Acceptance criteria

  • Window partition, cumulative lengths, permutation, and restoration match the existing MLX implementation exactly.
  • Full/window block selection is checkpoint-driven and artifact-fingerprinted.
  • Intermediate vision values and deterministic output match the pinned oracle within the documented stage tolerances/token exactness.
  • Multi-window and multi-media attention isolation is tested.
  • Image and advertised video behavior are either validated or explicitly rejected.
  • Qwen2-VL and text-only 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