lablup / lablup/mlxcel

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

Open
#866 0 comments 0 reactions 0 assignees View on GitHub
area:inference area:models priority:low status:review type:enhancement
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.