Fix Qwen3Omni VisionTransformer compatibility with updated vLLM Qwen2_5_VisionAttention

The upstream vLLM Qwen2_5_VisionAttention.forward() now requires a
`sequence_lengths` argument, but Qwen3_VisionBlock.forward() does not
pass it, causing a TypeError at init time.

Instead of patching the vLLM package directly, introduce a
Qwen3Omni_VisionTransformer subclass in vllm-omni that overrides
forward() to properly compute and propagate sequence_lengths via
MMEncoderAttention (following qwen3_vl.py's pattern). Both thinker
and talker now import the fixed subclass.

Also bump the vLLM wheel to e196028dbaa41b.

Made-with: Cursor

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

docker/Dockerfile.ci

@@ -18,7 +18,7 @@ RUN apt-get update && \
 # Note: --prerelease=allow is required because nightly/commit wheels use
 # pre-release version numbers (e.g. 2.10.0.dev...) which uv ignores by default,
 # causing it to fall back to the older stable release from PyPI.
-RUN uv pip install --system --upgrade --force-reinstall --prerelease=allow vllm --extra-index-url https://wheels.vllm.ai/3dd142adb7f7fa3b0b8345ee046ba36a12d9e530
+RUN uv pip install --system --upgrade --force-reinstall --prerelease=allow vllm --extra-index-url https://wheels.vllm.ai/e196028dbaa41b8e3d1c27b7dbd026403af76c27
 
 RUN uv pip install --system --no-cache-dir ".[dev]"
 RUN uv pip install --system --upgrade \

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py

@@ -19,7 +19,6 @@ from vllm.model_executor.models.interfaces import (
 from vllm.model_executor.models.qwen2_5_omni_thinker import (
     Qwen2_5OmniThinkerDummyInputsBuilder,
 )
-from vllm.model_executor.models.qwen3_omni_moe_thinker import Qwen3Omni_VisionTransformer
 from vllm.model_executor.models.utils import (
     AutoWeightsLoader,
     WeightsMapper,
@@ -33,6 +32,7 @@ from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_code_predictor_mt
 )
 from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import (
     Qwen3MoeLLMForCausalLM,
+    Qwen3Omni_VisionTransformer,
     Qwen3OmniMoeConditionalGenerationMixin,
     Qwen3OmniMoeThinkerMultiModalProcessor,
     Qwen3OmniMoeThinkerProcessingInfo,

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py

@@ -43,9 +43,12 @@ from transformers.models.qwen3_omni_moe.processing_qwen3_omni_moe import (
 from transformers.models.whisper import WhisperFeatureExtractor
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig
-from vllm.distributed import get_pp_group
+from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
 from vllm.inputs.data import PromptType
 from vllm.logger import init_logger
+from vllm.model_executor.layers.attention.mm_encoder_attention import (
+    MMEncoderAttention,
+)
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
 from vllm.model_executor.models.interfaces import (
@@ -70,7 +73,9 @@ from vllm.model_executor.models.qwen2_audio import Qwen2AudioProcessingInfo
 from vllm.model_executor.models.qwen3_moe import Qwen3MoeForCausalLM
 from vllm.model_executor.models.qwen3_moe import Qwen3MoeModel as _Qwen3MoeLLMModel
 from vllm.model_executor.models.qwen3_omni_moe_thinker import (
-    Qwen3Omni_VisionTransformer,
+    Qwen3Omni_VisionTransformer as _Qwen3Omni_VisionTransformer,
+)
+from vllm.model_executor.models.qwen3_omni_moe_thinker import (
     Qwen3OmniMoeAudioEncoder,
     _get_feat_extract_output_lengths,
 )
@@ -104,6 +109,90 @@ except (ImportError, ModuleNotFoundError):
 
 logger = init_logger(__name__)
 
+
+class Qwen3Omni_VisionTransformer(_Qwen3Omni_VisionTransformer):
+    """Subclass that fixes Qwen2_5_VisionAttention.forward() compatibility.
+
+    The upstream Qwen3_VisionBlock.forward() does not pass the
+    ``sequence_lengths`` argument required by the updated
+    Qwen2_5_VisionAttention.forward() signature.  This subclass overrides
+    ``forward()`` to compute ``sequence_lengths`` via MMEncoderAttention
+    (following the pattern in qwen3_vl.py) and calls block internals
+    directly so the argument is forwarded correctly.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.tp_size = get_tensor_model_parallel_world_size()
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        grid_thw,
+    ) -> torch.Tensor:
+        hidden_states = x.to(device=self.device, dtype=self.dtype)
+        hidden_states = self.patch_embed(hidden_states)
+
+        if self.apply_vit_abs_pos_embed:
+            pos_embeds = self.fast_pos_embed_interpolate(grid_thw)
+            hidden_states = hidden_states + pos_embeds
+        rotary_pos_emb_cos, rotary_pos_emb_sin = self.rot_pos_emb(grid_thw)
+
+        if isinstance(grid_thw, torch.Tensor):
+            grid_thw_np = grid_thw.cpu().numpy().astype(np.int32)
+        else:
+            grid_thw_np = np.array(grid_thw, dtype=np.int32)
+
+        cu_seqlens = np.repeat(grid_thw_np[:, 1] * grid_thw_np[:, 2], grid_thw_np[:, 0]).cumsum(axis=0, dtype=np.int32)
+        cu_seqlens = np.concatenate([np.zeros(1, dtype=np.int32), cu_seqlens])
+
+        sequence_lengths = MMEncoderAttention.maybe_compute_sequence_lengths(self.attn_backend, cu_seqlens)
+        if sequence_lengths is not None:
+            sequence_lengths = torch.from_numpy(sequence_lengths).to(self.device, non_blocking=True)
+        max_seqlen = torch.tensor(
+            MMEncoderAttention.compute_max_seqlen(self.attn_backend, cu_seqlens),
+            dtype=torch.int32,
+            device=self.device,
+        )
+        cu_seqlens = MMEncoderAttention.maybe_recompute_cu_seqlens(
+            self.attn_backend,
+            cu_seqlens,
+            self.hidden_size,
+            self.tp_size,
+        )
+        cu_seqlens = torch.from_numpy(cu_seqlens).to(self.device, non_blocking=True)
+
+        hidden_states = hidden_states.unsqueeze(1)
+
+        hidden_states_list = []
+        deepstack_visual_indexes = self.deepstack_visual_indexes
+
+        for layer_num, blk in enumerate(self.blocks):
+            hidden_states = hidden_states + blk.attn(
+                blk.norm1(hidden_states),
+                cu_seqlens=cu_seqlens,
+                rotary_pos_emb_cos=rotary_pos_emb_cos,
+                rotary_pos_emb_sin=rotary_pos_emb_sin,
+                max_seqlen=max_seqlen,
+                sequence_lengths=sequence_lengths,
+            )
+            hidden_states = hidden_states + blk.mlp(blk.norm2(hidden_states))
+
+            if deepstack_visual_indexes is not None and layer_num in deepstack_visual_indexes:
+                hidden_states_list.append(hidden_states)
+
+        hidden_states = self.merger(hidden_states)
+
+        if deepstack_visual_indexes is not None:
+            processed_hidden_states_list = [hidden_states]
+            for idx, x_ds in enumerate(hidden_states_list):
+                x_ds = self.merger_list[idx](x_ds)
+                processed_hidden_states_list.append(x_ds)
+            hidden_states = torch.cat(processed_hidden_states_list, dim=1)
+
+        return hidden_states
+
+
 # Speech input languages supported by Qwen3-Omni
 # From: https://huggingface.co/Qwen/Qwen3-Omni-30B-A3B-Instruct
 ISO639_1_SUPPORTED_LANGS = {