Enhance cu_seqlens computation in Qwen3Omni_VisionTransformer by switching to tensor operations for improved performance and compatibility. Added error handling to fallback on a vectorized implementation if necessary. This update aligns with upstream changes and optimizes the handling of sequence lengths during attention execution.

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py

@@ -29,6 +29,7 @@ from typing import Any, Literal, cast
 import numpy as np
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 from packaging.version import Version
 from transformers import PretrainedConfig
 from transformers import __version__ as TRANSFORMERS_VERSION
@@ -139,42 +140,67 @@ class Qwen3Omni_VisionTransformer(_Qwen3Omni_VisionTransformer):
         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)
+            grid_thw_tensor = grid_thw.to(self.device)
         else:
-            grid_thw_np = np.array(grid_thw, dtype=np.int32)
+            grid_thw_tensor = torch.as_tensor(grid_thw, dtype=torch.int32, device=self.device)
 
-        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_seq_lens(
-            self.attn_backend,
-            cu_seqlens,
-            self.device,
-        )
-        max_seqlen = torch.tensor(
-            MMEncoderAttention.compute_max_seqlen(self.attn_backend, cu_seqlens),
-            dtype=torch.int32,
-            device=self.device,
-        )
-        # Align with upstream vLLM qwen3_omni path: keep cu_seqlens in numpy for
-        # sequence length helpers and convert once for attention execution.
-        cu_seqlens = torch.from_numpy(cu_seqlens).to(self.device, non_blocking=True)
+        try:
+            cu_seqlens = torch.repeat_interleave(
+                grid_thw_tensor[:, 1] * grid_thw_tensor[:, 2],
+                grid_thw_tensor[:, 0],
+            ).cumsum(
+                dim=0,
+                dtype=grid_thw_tensor.dtype if torch.jit.is_tracing() else torch.int32,
+            )
+            cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)
+        except RuntimeError:
+            logger.warning(
+                "torch.repeat_interleave not executable, "
+                "switching to vectorized searchsorted implementation."
+            )
+            repeat_counts = grid_thw_tensor[:, 0]
+            values = grid_thw_tensor[:, 1] * grid_thw_tensor[:, 2]
+            repeat_cumsum = repeat_counts.cumsum(0)
+            total_items = repeat_cumsum[-1].item()
+            indices = torch.searchsorted(
+                repeat_cumsum,
+                torch.arange(total_items, device=grid_thw_tensor.device),
+                right=True,
+            )
+            cu_seqlens = values[indices].cumsum(
+                dim=0,
+                dtype=grid_thw_tensor.dtype if torch.jit.is_tracing() else torch.int32,
+            )
+            cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)
 
         hidden_states = hidden_states.unsqueeze(1)
+        rotary_pos_emb_cos = rotary_pos_emb_cos.to(hidden_states.device)
+        rotary_pos_emb_sin = rotary_pos_emb_sin.to(hidden_states.device)
+        max_seqlen = self.compute_attn_mask_seqlen(cu_seqlens)
+
+        grid_thw_np = grid_thw_tensor.cpu().numpy().astype(np.int32)
+        cu_seqlens_np = np.repeat(
+            grid_thw_np[:, 1] * grid_thw_np[:, 2], grid_thw_np[:, 0]
+        ).cumsum(axis=0, dtype=np.int32)
+        cu_seqlens_np = np.concatenate([np.zeros(1, dtype=np.int32), cu_seqlens_np])
+        sequence_lengths = MMEncoderAttention.maybe_compute_seq_lens(
+            self.attn_backend,
+            cu_seqlens_np,
+            self.device,
+        )
 
         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),
+            hidden_states = blk(
+                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)