feat(xla): support Qwen2.5-VL windowed vision attention
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
## Parent and dependency
- Parent epic: #566
- Blocked by #865
- Existing MLX implementation: `src/vision/qwen2_5_vl.rs` and `src/vision/encoders/qwen2_5_vl.rs`
## 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 #865 CLI/server/cancellation/accounting path and #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 #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.
Contributor guide
Research direction
Start with src/vision/qwen2_5_vl.rs and src/vision/encoders/qwen2_5_vl.rs, then review the qualified Qwen2-VL XLA path and the dependency issues #864 and #865. Compare processor tensors, window indices, cumulative lengths, reordered and restored features, and logits against the pinned MLX implementation and oracle; done requires the listed image/video behavior, attention-isolation tests, CLI/server regressions, and production-target validation.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- computer-vision, machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100