[VLM] Optimize get_rope_index for GLM4v (#17420)

Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>

benchmark/bench_rope/benchmark_rope_index.py

@@ -0,0 +1,425 @@
+# This script benchmarks MRotaryEmbedding.get_rope_index_glm4v (GLM4V mrope index builder).
+# It generates synthetic multimodal input_ids + attention_mask (+ optional image/video grids),
+# runs benchmarks.
+#
+# == Usage Examples ==
+#
+# python3 benchmark_rope_index.py --device cuda --num-tokens 1024 2048 --benchmark-iter 200
+
+import argparse
+import math
+import time
+from dataclasses import dataclass, field
+from typing import Any
+
+import numpy as np
+import torch
+
+from sglang.srt.layers.rotary_embedding import MRotaryEmbedding
+
+
+# -----------------------------
+# Minimal config objects
+# -----------------------------
+@dataclass
+class DummyVisionConfig:
+    spatial_merge_size: int = 2
+
+
+@dataclass
+class DummyHFConfig:
+    image_token_id: int = 32000
+    video_start_token_id: int = 32001
+    video_end_token_id: int = 32002
+    vision_config: DummyVisionConfig = field(
+        default_factory=lambda: DummyVisionConfig(spatial_merge_size=2)
+    )
+
+
+# -----------------------------
+# Helpers
+# -----------------------------
+def calculate_stats(times: list[float]) -> dict[str, float]:
+    """Calculate statistics from a list of times."""
+    times_array = np.array(times, dtype=np.float64)
+    return {
+        "mean": float(np.mean(times_array)),
+        "median": float(np.median(times_array)),
+        "p99": float(np.percentile(times_array, 99)),
+        "min": float(np.min(times_array)),
+        "max": float(np.max(times_array)),
+    }
+
+
+def _sync(device: torch.device):
+    if device.type == "cuda":
+        torch.cuda.synchronize()
+
+
+def _approx_hw(patches: int, merge: int) -> tuple[int, int]:
+    # want (h/merge)*(w/merge) ~= patches
+    gh = int(math.sqrt(max(1, patches)))
+    gw = max(1, patches // max(1, gh))
+    return gh * merge, gw * merge
+
+
+def generate_test_data(
+    num_tokens: int,
+    batch_size: int,
+    hf_config: DummyHFConfig,
+    dtype: torch.dtype,
+    device: torch.device,
+    pad_ratio: float,
+    num_images_per_sample: int,
+    image_patch_tokens: int,
+    num_videos_per_sample: int,
+    video_patch_tokens: int,
+    seed: int,
+):
+    """
+    Generate synthetic (input_ids, attention_mask, image_grid_thw, video_grid_thw).
+
+    NOTE:
+      - image_grid_thw / video_grid_thw are global lists across the entire batch in encounter order,
+        matching the function's image_index/video_index behavior.
+      - image patches are represented by repeated image_token_id.
+      - video patches are represented by image_token_id wrapped with start/end tokens.
+    """
+    torch.manual_seed(seed)
+
+    forbidden = {
+        0,
+        hf_config.image_token_id,
+        hf_config.video_start_token_id,
+        hf_config.video_end_token_id,
+    }
+    vocab_size = 50000
+
+    def rand_text(n: int) -> torch.Tensor:
+        # generate random ids not in forbidden
+        out = torch.randint(1, vocab_size, (n,), device=device, dtype=torch.long)
+        # fix forbidden by +1 until ok (cheap, deterministic enough for benchmark data)
+        for bad in forbidden:
+            out = torch.where(out == bad, out + 1, out)
+        return out
+
+    image_grids: list[list[int]] = []
+    video_grids: list[list[int]] = []
+
+    input_ids = torch.zeros((batch_size, num_tokens), device=device, dtype=torch.long)
+    attention_mask = torch.zeros(
+        (batch_size, num_tokens), device=device, dtype=torch.long
+    )
+
+    eff_len = int(round(num_tokens * (1.0 - pad_ratio)))
+    eff_len = max(1, min(num_tokens, eff_len))
+
+    min_needed = 1
+    min_needed += num_images_per_sample * image_patch_tokens
+    min_needed += num_videos_per_sample * (2 + video_patch_tokens)
+    if eff_len < min_needed:
+        num_images_per_sample = 0
+        num_videos_per_sample = 0
+
+    for b in range(batch_size):
+        blocks: list[torch.Tensor] = []
+
+        reserved = (
+            num_images_per_sample * image_patch_tokens
+            + num_videos_per_sample * (2 + video_patch_tokens)
+        )
+        reserved = min(reserved, max(0, eff_len - 1))
+        text_budget = max(1, eff_len - reserved)
+
+        n_text_chunks = num_images_per_sample + num_videos_per_sample + 1
+        base = text_budget // n_text_chunks
+        rem = text_budget % n_text_chunks
+        text_chunks = [base + (1 if i < rem else 0) for i in range(n_text_chunks)]
+
+        tci = 0
+        for _ in range(num_images_per_sample):
+            blocks.append(rand_text(text_chunks[tci]))
+            tci += 1
+            blocks.append(
+                torch.full(
+                    (image_patch_tokens,),
+                    hf_config.image_token_id,
+                    device=device,
+                    dtype=torch.long,
+                )
+            )
+
+            h, w = _approx_hw(
+                image_patch_tokens, hf_config.vision_config.spatial_merge_size
+            )
+            image_grids.append([1, h, w])
+
+        for _ in range(num_videos_per_sample):
+            blocks.append(rand_text(text_chunks[tci]))
+            tci += 1
+            blocks.append(
+                torch.tensor(
+                    [hf_config.video_start_token_id], device=device, dtype=torch.long
+                )
+            )
+            blocks.append(
+                torch.full(
+                    (video_patch_tokens,),
+                    hf_config.image_token_id,
+                    device=device,
+                    dtype=torch.long,
+                )
+            )
+            blocks.append(
+                torch.tensor(
+                    [hf_config.video_end_token_id], device=device, dtype=torch.long
+                )
+            )
+
+            h, w = _approx_hw(
+                video_patch_tokens, hf_config.vision_config.spatial_merge_size
+            )
+            # first field = group count used by code; set to 1
+            video_grids.append([1, h, w])
+
+        blocks.append(rand_text(text_chunks[tci]))
+
+        tokens = torch.cat(blocks, dim=0)[:eff_len]
+        pad = torch.zeros(
+            (num_tokens - tokens.numel(),), device=device, dtype=torch.long
+        )
+        ids = torch.cat([tokens, pad], dim=0)
+
+        mask = torch.cat(
+            [
+                torch.ones((tokens.numel(),), device=device, dtype=torch.long),
+                torch.zeros(
+                    (num_tokens - tokens.numel(),), device=device, dtype=torch.long
+                ),
+            ],
+            dim=0,
+        )
+
+        input_ids[b] = ids
+        attention_mask[b] = mask
+
+    image_grid_thw = (
+        torch.tensor(image_grids, device=device, dtype=torch.long)
+        if len(image_grids)
+        else None
+    )
+    video_grid_thw = (
+        torch.tensor(video_grids, device=device, dtype=torch.long)
+        if len(video_grids)
+        else None
+    )
+    return (
+        input_ids.to(dtype=torch.long),
+        attention_mask.to(dtype=torch.long),
+        image_grid_thw,
+        video_grid_thw,
+    )
+
+
+def benchmark_rope_index(
+    model_name: str,
+    tp_size: int,
+    num_tokens: int,
+    batch_size: int,
+    pad_ratio: float,
+    spatial_merge_size: int,
+    num_images: int,
+    image_patch_tokens: int,
+    num_videos: int,
+    video_patch_tokens: int,
+    dtype: torch.dtype,
+    seed: int,
+    warmup_iter: int,
+    benchmark_iter: int,
+    device: torch.device,
+):
+    torch.manual_seed(seed)
+    hf_config = DummyHFConfig(
+        image_token_id=32000,
+        video_start_token_id=32001,
+        video_end_token_id=32002,
+        vision_config=DummyVisionConfig(spatial_merge_size=spatial_merge_size),
+    )
+
+    print(80 * "=")
+    print(
+        f"Evaluating: {model_name} tp_size={tp_size} "
+        f"num_tokens={num_tokens} batch={batch_size} pad_ratio={pad_ratio} "
+        f"images/sample={num_images} image_patch_tokens={image_patch_tokens} "
+        f"videos/sample={num_videos} video_patch_tokens={video_patch_tokens} "
+        f"dtype={dtype} device={device}"
+    )
+
+    input_ids, attention_mask, image_grid_thw, video_grid_thw = generate_test_data(
+        num_tokens=num_tokens,
+        batch_size=batch_size,
+        hf_config=hf_config,
+        dtype=dtype,
+        device=device,
+        pad_ratio=pad_ratio,
+        num_images_per_sample=num_images,
+        image_patch_tokens=image_patch_tokens,
+        num_videos_per_sample=num_videos,
+        video_patch_tokens=video_patch_tokens,
+        seed=seed,
+    )
+
+    # Smoke test
+    has_mm = (image_grid_thw is not None) or (video_grid_thw is not None)
+    if has_mm:
+        pos, delta = MRotaryEmbedding.get_rope_index_glm4v(
+            input_ids=input_ids,
+            hf_config=hf_config,
+            image_grid_thw=image_grid_thw,
+            video_grid_thw=video_grid_thw,
+            attention_mask=attention_mask,
+        )
+        assert pos.shape == (3, batch_size, num_tokens)
+        assert delta.shape == (batch_size, 1)
+
+    # Warm up
+    for _ in range(warmup_iter):
+        if has_mm:
+            MRotaryEmbedding.get_rope_index_glm4v(
+                input_ids=input_ids,
+                hf_config=hf_config,
+                image_grid_thw=image_grid_thw,
+                video_grid_thw=video_grid_thw,
+                attention_mask=attention_mask,
+            )
+        MRotaryEmbedding.get_rope_index_glm4v(
+            input_ids=input_ids,
+            hf_config=hf_config,
+            image_grid_thw=None,
+            video_grid_thw=None,
+            attention_mask=attention_mask,
+        )
+
+    _sync(device)
+
+    # Time multimodal branch
+    multimodal_times = []
+    for _ in range(benchmark_iter):
+        _sync(device)
+        start = time.time()
+        MRotaryEmbedding.get_rope_index_glm4v(
+            input_ids=input_ids,
+            hf_config=hf_config,
+            image_grid_thw=image_grid_thw,
+            video_grid_thw=video_grid_thw,
+            attention_mask=attention_mask,
+        )
+        _sync(device)
+        multimodal_times.append(time.time() - start)
+
+    # Time fallback branch
+    fallback_times = []
+    for _ in range(benchmark_iter):
+        _sync(device)
+        start = time.time()
+        MRotaryEmbedding.get_rope_index_glm4v(
+            input_ids=input_ids,
+            hf_config=hf_config,
+            image_grid_thw=None,
+            video_grid_thw=None,
+            attention_mask=attention_mask,
+        )
+        _sync(device)
+        fallback_times.append(time.time() - start)
+
+    multimodal_stats = calculate_stats(multimodal_times)
+    fallback_stats = calculate_stats(fallback_times)
+
+    print(f"\nPerformance for config (B={batch_size}, T={num_tokens}):")
+    print(
+        f"Multimodal: mean={multimodal_stats['mean']:.8f}s, "
+        f"median={multimodal_stats['median']:.8f}s, "
+        f"p99={multimodal_stats['p99']:.8f}s"
+    )
+    print(
+        f"Fallback:   mean={fallback_stats['mean']:.8f}s, "
+        f"median={fallback_stats['median']:.8f}s, "
+        f"p99={fallback_stats['p99']:.8f}s"
+    )
+
+    if has_mm:
+        speedup = (
+            multimodal_stats["mean"] / fallback_stats["mean"]
+            if fallback_stats["mean"] > 0
+            else float("inf")
+        )
+        print(f"Fallback Speedup over Multimodal: {speedup:.8f}x")
+    else:
+        speedup = float("nan")
+        print(
+            "[INFO] num_tokens too small for multimodal segments; skip multimodal benchmark."
+        )
+
+    print(f"Fallback Speedup over Multimodal: {speedup:.8f}x")
+
+    return multimodal_stats, fallback_stats, speedup
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Benchmark GLM4V get_rope_index_glm4v."
+    )
+    parser.add_argument("--model-name", type=str, default="GLM4V")
+    parser.add_argument("--tp-size", type=int, default=1)
+    parser.add_argument(
+        "--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu"
+    )
+    parser.add_argument("--warmup-iter", type=int, default=10)
+    parser.add_argument("--benchmark-iter", type=int, default=100)
+    parser.add_argument("--dtype", type=str, choices=["int64"], default="int64")
+    parser.add_argument("--seed", type=int, default=0)
+
+    # token length sweep
+    parser.add_argument("--num-tokens", type=int, nargs="+", required=False)
+
+    # data shape knobs
+    parser.add_argument("--batch-size", type=int, default=1)
+    parser.add_argument("--pad-ratio", type=float, default=0.0)
+    parser.add_argument("--spatial-merge-size", type=int, default=2)
+    parser.add_argument("--num-images", type=int, default=1)
+    parser.add_argument("--image-patch-tokens", type=int, default=256)
+    parser.add_argument("--num-videos", type=int, default=1)
+    parser.add_argument("--video-patch-tokens", type=int, default=256)
+
+    # output
+    parser.add_argument("--out-dir", type=str, default=".")
+    args = parser.parse_args()
+    print(args)
+
+    device = torch.device(args.device)
+
+    if args.num_tokens is None:
+        num_tokens_list = [2**i for i in range(0, 18)]
+    else:
+        num_tokens_list = args.num_tokens
+
+    rows: list[dict[str, Any]] = []
+
+    for num_tokens in num_tokens_list:
+        multimodal_stats, fallback_stats, speedup = benchmark_rope_index(
+            model_name=args.model_name,
+            tp_size=args.tp_size,
+            num_tokens=num_tokens,
+            batch_size=args.batch_size,
+            pad_ratio=args.pad_ratio,
+            spatial_merge_size=args.spatial_merge_size,
+            num_images=args.num_images,
+            image_patch_tokens=args.image_patch_tokens,
+            num_videos=args.num_videos,
+            video_patch_tokens=args.video_patch_tokens,
+            dtype=getattr(torch, args.dtype),
+            seed=args.seed,
+            warmup_iter=args.warmup_iter,
+            benchmark_iter=args.benchmark_iter,
+            device=device,
+        )

python/sglang/srt/layers/rotary_embedding.py

@@ -1700,11 +1700,12 @@ class MRotaryEmbedding(RotaryEmbedding):
                         video_index += 1
                         remain_videos -= 1
                         ed = ed_video
-                    llm_grid_t, llm_grid_h, llm_grid_w = (
+                    # Avoid .item() lookups in repeated context
-                        t.item(),
+                    t_int, h_int, w_int = int(t), int(h), int(w)
-                        h.item() // spatial_merge_size,
+
-                        w.item() // spatial_merge_size,
+                    llm_grid_t = t_int
-                    )
+                    llm_grid_h = h_int // spatial_merge_size
+                    llm_grid_w = w_int // spatial_merge_size
                     text_len = ed - st
 
                     st_idx = (
@@ -1737,24 +1738,24 @@ class MRotaryEmbedding(RotaryEmbedding):
                         "qwen3_vl_moe",
                     ):
                         t_index = (
-                            torch.arange(llm_grid_t)
+                            torch.arange(llm_grid_t, device=position_ids.device)
                             .view(-1, 1)
-                            .expand(-1, llm_grid_h * llm_grid_w)
+                            .expand(llm_grid_t, llm_grid_h * llm_grid_w)
-                            .flatten()
+                            .reshape(-1)
                         )
                     else:
                         raise RuntimeError(f"Unimplemented model type: {model_type}")
                     h_index = (
-                        torch.arange(llm_grid_h)
+                        torch.arange(llm_grid_h, device=position_ids.device)
                         .view(1, -1, 1)
-                        .expand(llm_grid_t, -1, llm_grid_w)
+                        .expand(llm_grid_t, llm_grid_h, llm_grid_w)
-                        .flatten()
+                        .reshape(-1)
                     )
                     w_index = (
-                        torch.arange(llm_grid_w)
+                        torch.arange(llm_grid_w, device=position_ids.device)
                         .view(1, 1, -1)
-                        .expand(llm_grid_t, llm_grid_h, -1)
+                        .expand(llm_grid_t, llm_grid_h, llm_grid_w)
-                        .flatten()
+                        .reshape(-1)
                     )
                     llm_pos_ids_list.append(
                         torch.stack([t_index, h_index, w_index]) + text_len + st_idx
@@ -1787,10 +1788,9 @@ class MRotaryEmbedding(RotaryEmbedding):
             position_ids = (
                 position_ids.unsqueeze(0).expand(3, -1, -1).to(input_ids.device)
             )
-            max_position_ids = position_ids.max(0, keepdim=False)[0].max(
+            max_position_ids = position_ids.amax(dim=0, keepdim=False)
-                -1, keepdim=True
+            mrope_position_deltas = max_position_ids.amax(-1, keepdim=True) + 1 - s
-            )[0]
+
-            mrope_position_deltas = max_position_ids + 1 - s
             return position_ids, mrope_position_deltas
 
     @staticmethod
@@ -2107,13 +2107,17 @@ class MRotaryEmbedding(RotaryEmbedding):
         video_end_token_id = hf_config.video_end_token_id
         spatial_merge_size = hf_config.vision_config.spatial_merge_size
 
+        # Preallocate lists for efficiency
         mrope_position_deltas = []
+
         if input_ids is not None and (
             image_grid_thw is not None or video_grid_thw is not None
         ):
             total_input_ids = input_ids
+
             if attention_mask is None:
                 attention_mask = torch.ones_like(total_input_ids)
+
             position_ids = torch.ones(
                 3,
                 input_ids.shape[0],
@@ -2121,28 +2125,36 @@ class MRotaryEmbedding(RotaryEmbedding):
                 dtype=input_ids.dtype,
                 device=input_ids.device,
             )
+
             image_index, video_index = 0, 0
             video_group_index = 0
+            # Move attention mask to device once to avoid repeated transfers
             attention_mask = attention_mask.to(total_input_ids.device)
-            for i, input_ids in enumerate(total_input_ids):
-                input_ids = input_ids[attention_mask[i] == 1]
-                input_tokens = input_ids.tolist()
 
-                input_token_type = []
+            for i, ids in enumerate(total_input_ids):
+                curr_mask = attention_mask[i]
+                ids_masked = ids[curr_mask == 1]
+
+                # Preallocate input_token_type for maximum speed
+                input_tokens = ids_masked.tolist()
+                input_token_type = [""] * len(input_tokens)
+
+                # Single pass through tokens for type assignment, using explicit indices for performance
                 video_check_flg = False
-                for token in input_tokens:
+                for j, token in enumerate(input_tokens):
                     if token == video_start_token_id:
                         video_check_flg = True
                     elif token == video_end_token_id:
                         video_check_flg = False
 
                     if token == image_token_id and not video_check_flg:
-                        input_token_type.append("image")
+                        input_token_type[j] = "image"
                     elif token == image_token_id and video_check_flg:
-                        input_token_type.append("video")
+                        input_token_type[j] = "video"
                     else:
-                        input_token_type.append("text")
+                        input_token_type[j] = "text"
 
+                # Use itertools.groupby for consecutive token type groups (unchanged logic)
                 input_type_group = []
                 for key, group in itertools.groupby(
                     enumerate(input_token_type), lambda x: x[1]
@@ -2154,12 +2166,13 @@ class MRotaryEmbedding(RotaryEmbedding):
 
                 llm_pos_ids_list = []
                 video_frame_num = 1
+
                 for modality_type, start_idx, end_idx in input_type_group:
-                    st_idx = (
+                    # st_idx can be computed by torch directly for speed
-                        llm_pos_ids_list[-1].max() + 1
+                    if llm_pos_ids_list:
-                        if len(llm_pos_ids_list) > 0
+                        st_idx = llm_pos_ids_list[-1].max().item() + 1
-                        else 0
+                    else:
-                    )
+                        st_idx = 0
 
                     if modality_type == "image":
                         t, h, w = (
@@ -2167,103 +2180,102 @@ class MRotaryEmbedding(RotaryEmbedding):
                             image_grid_thw[image_index][1],
                             image_grid_thw[image_index][2],
                         )
-                        llm_grid_t, llm_grid_h, llm_grid_w = (
+                        # Avoid .item() lookups in repeated context
-                            t.item(),
+                        t_int, h_int, w_int = int(t), int(h), int(w)
-                            h.item() // spatial_merge_size,
-                            w.item() // spatial_merge_size,
-                        )
 
+                        llm_grid_t = t_int
+                        llm_grid_h = h_int // spatial_merge_size
+                        llm_grid_w = w_int // spatial_merge_size
+
+                        # Avoid unnecessary views/expands for speed, always flatten at the end
                         t_index = (
-                            torch.arange(llm_grid_t)
+                            torch.arange(llm_grid_t, device=position_ids.device)
                             .view(-1, 1)
-                            .expand(-1, llm_grid_h * llm_grid_w)
+                            .expand(llm_grid_t, llm_grid_h * llm_grid_w)
-                            .flatten()
+                            .reshape(-1)
                         )
                         h_index = (
-                            torch.arange(llm_grid_h)
+                            torch.arange(llm_grid_h, device=position_ids.device)
                             .view(1, -1, 1)
-                            .expand(llm_grid_t, -1, llm_grid_w)
+                            .expand(llm_grid_t, llm_grid_h, llm_grid_w)
-                            .flatten()
+                            .reshape(-1)
                         )
                         w_index = (
-                            torch.arange(llm_grid_w)
+                            torch.arange(llm_grid_w, device=position_ids.device)
                             .view(1, 1, -1)
-                            .expand(llm_grid_t, llm_grid_h, -1)
+                            .expand(llm_grid_t, llm_grid_h, llm_grid_w)
-                            .flatten()
+                            .reshape(-1)
                         )
                         llm_pos_ids_list.append(
                             torch.stack([t_index, h_index, w_index]) + st_idx
                         )
-
                         image_index += 1
                         video_frame_num = 1
 
                     elif modality_type == "video":
-                        t, h, w = (
+                        t = video_frame_num
-                            video_frame_num,
+                        h = video_grid_thw[video_index][1]
-                            video_grid_thw[video_index][1],
+                        w = video_grid_thw[video_index][2]
-                            video_grid_thw[video_index][2],
-                        )
 
-                        llm_grid_t, llm_grid_h, llm_grid_w = (
+                        h_int, w_int = int(h), int(w)
-                            t,
+                        llm_grid_h = h_int // spatial_merge_size
-                            h.item() // spatial_merge_size,
+                        llm_grid_w = w_int // spatial_merge_size
-                            w.item() // spatial_merge_size,
-                        )
 
-                        for t_idx in range(llm_grid_t):
+                        # Only one video frame at a time
+                        for t_idx in range(t):
                             t_index = (
-                                torch.tensor(t_idx)
+                                torch.tensor(t_idx, device=position_ids.device)
                                 .view(-1, 1)
-                                .expand(-1, llm_grid_h * llm_grid_w)
+                                .expand(1, llm_grid_h * llm_grid_w)
-                                .flatten()
+                                .reshape(-1)
                             )
-
                             h_index = (
-                                torch.arange(llm_grid_h)
+                                torch.arange(llm_grid_h, device=position_ids.device)
                                 .view(1, -1, 1)
-                                .expand(1, -1, llm_grid_w)
+                                .expand(1, llm_grid_h, llm_grid_w)
-                                .flatten()
+                                .reshape(-1)
                             )
                             w_index = (
-                                torch.arange(llm_grid_w)
+                                torch.arange(llm_grid_w, device=position_ids.device)
                                 .view(1, 1, -1)
-                                .expand(1, llm_grid_h, -1)
+                                .expand(1, llm_grid_h, llm_grid_w)
-                                .flatten()
+                                .reshape(-1)
                             )
                             llm_pos_ids_list.append(
                                 torch.stack([t_index, h_index, w_index]) + st_idx
                             )
 
                         video_group_index += 1
-
                         if video_group_index >= video_grid_thw[video_index][0]:
                             video_index += 1
                             video_group_index = 0
 
                         video_frame_num += 1
 
-                    else:
+                    else:  # text
                         text_len = end_idx - start_idx
-                        llm_pos_ids_list.append(
+                        # Use in-place expand for improved performance
-                            torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx
+                        text_range = torch.arange(text_len, device=position_ids.device)
-                        )
+                        text_pos = text_range.view(1, -1).expand(3, text_len) + st_idx
-
+                        llm_pos_ids_list.append(text_pos)
                         video_frame_num = 1
 
+                # Concatenate once outside for speed
                 llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
-                position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(
+                # Use advanced indexing for assignment
-                    position_ids.device
+                idx_mask = curr_mask == 1
-                )
+                position_ids[..., i, idx_mask] = llm_positions.to(position_ids.device)
                 mrope_position_deltas.append(
                     llm_positions.max() + 1 - len(total_input_ids[i])
                 )
+            # Build tensor in one call at the end
             mrope_position_deltas = torch.tensor(
                 mrope_position_deltas, device=input_ids.device
             ).unsqueeze(1)
             return position_ids, mrope_position_deltas
         else:
             if attention_mask is not None:
+                # Use in-place operations whenever possible
                 position_ids = attention_mask.long().cumsum(-1) - 1
                 position_ids.masked_fill_(attention_mask == 0, 1)
                 position_ids = (
@@ -2271,22 +2283,25 @@ class MRotaryEmbedding(RotaryEmbedding):
                     .expand(3, -1, -1)
                     .to(attention_mask.device)
                 )
-                max_position_ids = position_ids.max(0, keepdim=False)[0].max(
+                max_position_ids = position_ids.amax(dim=0, keepdim=False)
-                    -1, keepdim=True
+                mrope_position_deltas = (
-                )[0]
+                    max_position_ids.amax(-1, keepdim=True)
-                mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1]
+                    + 1
-            else:
+                    - attention_mask.shape[-1]
-                position_ids = (
-                    torch.arange(input_ids.shape[1], device=input_ids.device)
-                    .view(1, 1, -1)
-                    .expand(3, input_ids.shape[0], -1)
                 )
+            else:
+                length = input_ids.shape[1]
+                batch_size = input_ids.shape[0]
+                # Use torch.arange with in-place expansion
+                arange_ids = torch.arange(length, device=input_ids.device).view(
+                    1, 1, -1
+                )
+                position_ids = arange_ids.expand(3, batch_size, length)
                 mrope_position_deltas = torch.zeros(
-                    [input_ids.shape[0], 1],
+                    [batch_size, 1],
                     device=input_ids.device,
                     dtype=input_ids.dtype,
                 )
-
             return position_ids, mrope_position_deltas
 
     @staticmethod